From 788a0c35a9b55d5383a86938c3c0e9e5724f82ff Mon Sep 17 00:00:00 2001 From: yongman Date: Thu, 25 Jun 2026 20:07:00 +0800 Subject: [PATCH 1/7] LAC: fallback to legacy group when pd not support keyspace Signed-off-by: yongman --- .../LocalAdmissionController.cpp | 88 +++++++++-- .../LocalAdmissionController.h | 99 ++++++++++-- .../gtest_local_admission_controller.cpp | 147 ++++++++++++++++++ 3 files changed, 309 insertions(+), 25 deletions(-) create mode 100644 dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp index 359cdf157c8..8922740d677 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp @@ -59,6 +59,9 @@ uint64_t ResourceGroup::getPriority(uint64_t max_ru_per_sec) const std::optional ResourceGroup::buildRequestInfoIfNecessary(const SteadyClock::time_point & now) { + if (!enable_gac) + return {}; + std::lock_guard lock(mu); if (!beginRequestWithoutLock(now)) { @@ -292,6 +295,24 @@ LACRUConsumptionDeltaInfo ResourceGroup::updateRUConsumptionDeltaInfoWithoutLock return info; } +resource_manager::ResourceGroup LocalAdmissionController::buildReservedDefaultResourceGroup(const KeyspaceID & keyspace_id) +{ + resource_manager::ResourceGroup group_pb; + group_pb.set_name("default"); + group_pb.set_mode(resource_manager::GroupMode::RUMode); + group_pb.set_priority(ResourceGroup::UserMediumPriority); + group_pb.mutable_keyspace_id()->set_value(keyspace_id); + auto * settings = group_pb.mutable_r_u_settings()->mutable_r_u()->mutable_settings(); + settings->set_fill_rate(RESERVED_DEFAULT_RESOURCE_GROUP_RU_PER_SEC); + settings->set_burst_limit(-1); + return group_pb; +} + +void LocalAdmissionController::addReservedDefaultResourceGroup(const KeyspaceID & keyspace_id) +{ + addResourceGroup(keyspace_id, buildReservedDefaultResourceGroup(keyspace_id), false); +} + void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & keyspace_id, const std::string & name) { if (unlikely(stopped)) @@ -300,9 +321,11 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k if (name.empty()) return; - ResourceGroupPtr group = findResourceGroup(keyspace_id, name); - if (group != nullptr) - return; + { + std::lock_guard lock(mu); + if (findResourceGroupWithCompatWithoutLock(keyspace_id, name).group != nullptr) + return; + } resource_manager::GetResourceGroupRequest req; req.mutable_keyspace_id()->set_value(keyspace_id); @@ -323,12 +346,54 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k getCurrentExceptionMessage(false)); } - RUNTIME_CHECK_MSG( - !resp.has_error(), - "warmupResourceGroupInfoCache: {}(keyspace={}) failed: {}", - name, - keyspace_id, - resp.error().message()); + if (resp.has_error()) + { + resource_manager::GetResourceGroupRequest legacy_req; + legacy_req.set_resource_group_name(name); + resource_manager::GetResourceGroupResponse legacy_resp; + try + { + legacy_resp = cluster->pd_client->getResourceGroup(legacy_req); + } + catch (...) + { + LOG_WARNING( + log, + "warmupResourceGroupInfoCache failed to fetch legacy rg for keyspace={}, rg={}, err={}", + keyspace_id, + name, + getCurrentExceptionMessage(false)); + } + if (!legacy_resp.has_error()) + { + checkGACRespValid(legacy_resp.group()); + const auto rg_keyspace_id + = legacy_resp.group().has_keyspace_id() ? legacy_resp.group().keyspace_id().value() : NullspaceID; + addResourceGroup(rg_keyspace_id, legacy_resp.group()); + return; + } + if (isReservedDefaultResourceGroup(name)) + { + { + std::lock_guard lock(mu); + if (findResourceGroupWithoutLock(NullspaceID, name) != nullptr) + return; + } + LOG_WARNING( + log, + "warmupResourceGroupInfoCache fallback to reserved default rg for keyspace={}, err={}", + keyspace_id, + resp.error().message()); + addReservedDefaultResourceGroup(keyspace_id); + return; + } + RUNTIME_CHECK_MSG( + false, + "warmupResourceGroupInfoCache: {}(keyspace={}) failed: {}", + name, + keyspace_id, + resp.error().message()); + } checkGACRespValid(resp.group()); @@ -470,7 +535,8 @@ std::optional LocalAdmissionController::b for (const auto & info : request_infos) { auto * group_request = gac_req.add_requests(); - group_request->mutable_keyspace_id()->set_value(info.keyspace_id); + if (info.keyspace_id != NullspaceID) + group_request->mutable_keyspace_id()->set_value(info.keyspace_id); group_request->set_resource_group_name(info.resource_group_name); assert(info.acquire_tokens > 0.0 || info.ru_consumption_delta > 0.0 || is_final_report); if (info.acquire_tokens > 0.0 || is_final_report) @@ -535,7 +601,7 @@ static std::vector> extractGACReqNames( std::vector> res; res.reserve(gac_req.requests_size()); for (const auto & req : gac_req.requests()) - res.push_back({req.keyspace_id().value(), req.resource_group_name()}); + res.push_back({req.has_keyspace_id() ? req.keyspace_id().value() : NullspaceID, req.resource_group_name()}); return res; } diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h index 6bd1cdea4d6..bcea1123c60 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h @@ -46,6 +46,12 @@ struct GACRequestInfo double ru_consumption_delta; }; +struct ResourceGroupCompatLookupResult +{ + std::shared_ptr group; + bool is_fallback = false; +}; + inline std::string getResourceGroupMetricName(const KeyspaceID & keyspace_id, const std::string & name) { return fmt::format("{}-{}", keyspace_id, name); @@ -74,10 +80,12 @@ class ResourceGroup final : private boost::noncopyable explicit ResourceGroup( const KeyspaceID & keyspace_id_, const resource_manager::ResourceGroup & group_pb_, - const SteadyClock::time_point & tp) + const SteadyClock::time_point & tp, + bool enable_gac_ = true) : keyspace_id(keyspace_id_) , name(group_pb_.name()) , name_with_keyspace_id(getResourceGroupMetricName(keyspace_id_, group_pb_.name())) + , enable_gac(enable_gac_) , group_pb(group_pb_) , log(Logger::get("resource group:" + group_pb_.name())) { @@ -299,6 +307,7 @@ class ResourceGroup final : private boost::noncopyable const KeyspaceID keyspace_id; const std::string name; const std::string name_with_keyspace_id; // For metrics. + bool enable_gac = true; uint32_t user_priority_val = 0; uint64_t user_ru_per_sec = 0; bool burstable = false; @@ -355,17 +364,29 @@ struct LACPairHash class LocalAdmissionController final : private boost::noncopyable { public: - LocalAdmissionController(::pingcap::kv::Cluster * cluster_, Etcd::ClientPtr etcd_client_, bool with_keyspace) + LocalAdmissionController( + ::pingcap::kv::Cluster * cluster_, + Etcd::ClientPtr etcd_client_, + bool with_keyspace, + bool start_background_threads = true) : cluster(cluster_) , etcd_client(etcd_client_) , watch_gac_grpc_context(std::make_unique()) { - background_threads.emplace_back([this] { this->mainLoop(); }); - background_threads.emplace_back([this] { this->requestGACLoop(); }); - if (with_keyspace) - background_threads.emplace_back([this] { this->watchGACLoop(GAC_KEYSPACE_RESOURCE_GROUP_ETCD_PATH); }); - else - background_threads.emplace_back([this] { this->watchGACLoop(GAC_RESOURCE_GROUP_ETCD_PATH); }); + if (start_background_threads) + { + background_threads.emplace_back([this] { this->mainLoop(); }); + background_threads.emplace_back([this] { this->requestGACLoop(); }); + if (with_keyspace) + { + background_threads.emplace_back([this] { this->watchGACLoop(GAC_KEYSPACE_RESOURCE_GROUP_ETCD_PATH); }); + background_threads.emplace_back([this] { this->watchGACLoop(GAC_RESOURCE_GROUP_ETCD_PATH); }); + } + else + { + background_threads.emplace_back([this] { this->watchGACLoop(GAC_RESOURCE_GROUP_ETCD_PATH); }); + } + } current_tick = SteadyClock::now(); last_clear_cpu_time = current_tick; @@ -463,6 +484,20 @@ class LocalAdmissionController final : private boost::noncopyable refill_token_callback = nullptr; } +#ifdef DBMS_PUBLIC_GTEST + size_t cachedResourceGroupCount() const + { + std::lock_guard lock(mu); + return keyspace_resource_groups.size(); + } + + ResourceGroupPtr getCachedResourceGroupForTest(const KeyspaceID & keyspace_id, const std::string & name) const + { + std::lock_guard lock(mu); + return findResourceGroupWithoutLock(keyspace_id, name); + } +#endif + #ifdef DBMS_PUBLIC_GTEST static std::unique_ptr global_instance; #else @@ -476,6 +511,7 @@ class LocalAdmissionController final : private boost::noncopyable static constexpr auto DEGRADE_MODE_DURATION = std::chrono::seconds(120); static constexpr double ACQUIRE_RU_AMPLIFICATION = 1.1; static constexpr auto DEFAULT_MAX_EST_WAIT_DURATION = std::chrono::milliseconds(1000); + static constexpr uint64_t RESERVED_DEFAULT_RESOURCE_GROUP_RU_PER_SEC = std::numeric_limits::max(); #ifndef DBMS_PUBLIC_GTEST private: @@ -519,7 +555,7 @@ class LocalAdmissionController final : private boost::noncopyable { { std::lock_guard lock(mu); - keyspace_low_token_resource_groups.insert({keyspace_id, name}); + keyspace_low_token_resource_groups.insert({group->getKeyspaceID(), group->getName()}); } cv.notify_all(); } @@ -531,14 +567,14 @@ class LocalAdmissionController final : private boost::noncopyable ResourceGroupPtr findResourceGroup(const KeyspaceID & keyspace_id, const std::string & name) const { std::lock_guard lock(mu); - return findResourceGroupWithoutLock(keyspace_id, name); + return findResourceGroupWithCompatWithoutLock(keyspace_id, name).group; } std::pair findResourceGroupAndMaxRUPerSec( const KeyspaceID & keyspace_id, const std::string & name) const { std::lock_guard lock(mu); - auto rg = findResourceGroupWithoutLock(keyspace_id, name); + auto rg = findResourceGroupWithCompatWithoutLock(keyspace_id, name).group; return {rg, max_ru_per_sec}; } ResourceGroupPtr findResourceGroupWithoutLock(const KeyspaceID & keyspace_id, const std::string & name) const @@ -548,8 +584,29 @@ class LocalAdmissionController final : private boost::noncopyable return nullptr; return iter->second; } + ResourceGroupCompatLookupResult findResourceGroupWithCompatWithoutLock( + const KeyspaceID & keyspace_id, + const std::string & name) const + { + if (auto group = findResourceGroupWithoutLock(keyspace_id, name); likely(group != nullptr)) + return {.group = group, .is_fallback = false}; + + if (keyspace_id != NullspaceID) + { + if (auto group = findResourceGroupWithoutLock(NullspaceID, name); likely(group != nullptr)) + return {.group = group, .is_fallback = true}; + } + return {}; + } void addResourceGroup(const KeyspaceID & keyspace_id, const resource_manager::ResourceGroup & new_group_pb) + { + addResourceGroup(keyspace_id, new_group_pb, true); + } + void addResourceGroup( + const KeyspaceID & keyspace_id, + const resource_manager::ResourceGroup & new_group_pb, + bool enable_gac) { uint64_t user_ru_per_sec = new_group_pb.r_u_settings().r_u().settings().fill_rate(); std::lock_guard lock(mu); @@ -559,10 +616,21 @@ class LocalAdmissionController final : private boost::noncopyable auto iter = keyspace_resource_groups.find({keyspace_id, new_group_pb.name()}); if (iter != keyspace_resource_groups.end()) - return; + { + if (iter->second->enable_gac || !enable_gac) + return; + + updateMaxRUPerSecAfterDeleteWithoutLock(iter->second->user_ru_per_sec); + keyspace_resource_groups.erase(iter); + } - LOG_INFO(log, "add new resource group, info: {}", new_group_pb.ShortDebugString()); - auto new_group = std::make_shared(keyspace_id, new_group_pb, current_tick); + LOG_INFO( + log, + "add new resource group, keyspace={} enable_gac={} info: {}", + keyspace_id, + enable_gac, + new_group_pb.ShortDebugString()); + auto new_group = std::make_shared(keyspace_id, new_group_pb, current_tick, enable_gac); keyspace_resource_groups.insert({{keyspace_id, new_group_pb.name()}, new_group}); if (refill_token_callback) @@ -600,6 +668,9 @@ class LocalAdmissionController final : private boost::noncopyable // requestGACLoop related methods. void doRequestGAC(); + static bool isReservedDefaultResourceGroup(const std::string & name) { return name == "default"; } + static resource_manager::ResourceGroup buildReservedDefaultResourceGroup(const KeyspaceID & keyspace_id); + void addReservedDefaultResourceGroup(const KeyspaceID & keyspace_id); // watchGACLoop related methods. void doWatch(const std::string & etcd_path); diff --git a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp new file mode 100644 index 00000000000..ecc36806b95 --- /dev/null +++ b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp @@ -0,0 +1,147 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed 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. + +#include +#include +#include +#include + +namespace DB::tests +{ +namespace +{ +class TestPDClient final : public pingcap::pd::MockPDClient +{ +public: + explicit TestPDClient( + std::function + get_resource_group_) + : get_resource_group(std::move(get_resource_group_)) + {} + + resource_manager::GetResourceGroupResponse getResourceGroup( + const resource_manager::GetResourceGroupRequest & req) override + { + return get_resource_group(req); + } + + resource_manager::TokenBucketsResponse acquireTokenBuckets(const resource_manager::TokenBucketsRequest &) override + { + return {}; + } + +private: + std::function + get_resource_group; +}; +} // namespace + +TEST(LocalAdmissionControllerTest, ReservedDefaultGroupFallsBackToLocalCacheWhenPDDoesNotPersistIt) +{ + constexpr KeyspaceID keyspace_id = 7; + pingcap::kv::Cluster cluster; + cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest &) { + resource_manager::GetResourceGroupResponse resp; + resp.mutable_error()->set_message("resource group default not found"); + return resp; + }); + + LocalAdmissionController lac(&cluster, nullptr, false, false); + + ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); + ASSERT_EQ(lac.cachedResourceGroupCount(), 1); + + auto priority = lac.getPriority(keyspace_id, "default"); + ASSERT_TRUE(priority.has_value()); + + auto group = lac.getCachedResourceGroupForTest(keyspace_id, "default"); + ASSERT_NE(group, nullptr); + EXPECT_FALSE(group->enable_gac); + EXPECT_EQ(group->user_ru_per_sec, LocalAdmissionController::RESERVED_DEFAULT_RESOURCE_GROUP_RU_PER_SEC); + EXPECT_EQ(group->group_pb.priority(), ResourceGroup::UserMediumPriority); + EXPECT_TRUE(group->group_pb.has_keyspace_id()); + EXPECT_EQ(group->group_pb.keyspace_id().value(), keyspace_id); + EXPECT_FALSE(group->buildRequestInfoIfNecessary(SteadyClock::now()).has_value()); +} + +TEST(LocalAdmissionControllerTest, ReservedDefaultGroupUsesPDConfigWhenPresent) +{ + constexpr KeyspaceID keyspace_id = 9; + pingcap::kv::Cluster cluster; + cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { + resource_manager::GetResourceGroupResponse resp; + auto * group = resp.mutable_group(); + group->set_name(req.resource_group_name()); + group->set_mode(resource_manager::GroupMode::RUMode); + group->set_priority(ResourceGroup::UserLowPriority); + group->mutable_keyspace_id()->set_value(req.keyspace_id().value()); + auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); + settings->set_fill_rate(1024); + settings->set_burst_limit(1024); + return resp; + }); + + LocalAdmissionController lac(&cluster, nullptr, false, false); + + ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); + + auto group = lac.getCachedResourceGroupForTest(keyspace_id, "default"); + ASSERT_NE(group, nullptr); + EXPECT_TRUE(group->enable_gac); + EXPECT_EQ(group->user_ru_per_sec, 1024); + EXPECT_EQ(group->group_pb.priority(), ResourceGroup::UserLowPriority); + + auto priority = lac.getPriority(keyspace_id, "default"); + ASSERT_TRUE(priority.has_value()); +} + +TEST(LocalAdmissionControllerTest, KeyspaceLookupFallsBackToLegacyCachedGroup) +{ + constexpr KeyspaceID keyspace_id = 11; + pingcap::kv::Cluster cluster; + cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { + resource_manager::GetResourceGroupResponse resp; + if (req.has_keyspace_id()) + { + resp.mutable_error()->set_message("resource group default not found"); + return resp; + } + + auto * group = resp.mutable_group(); + group->set_name(req.resource_group_name()); + group->set_mode(resource_manager::GroupMode::RUMode); + group->set_priority(ResourceGroup::UserHighPriority); + auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); + settings->set_fill_rate(2048); + settings->set_burst_limit(2048); + return resp; + }); + + LocalAdmissionController lac(&cluster, nullptr, false, false); + + ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); + ASSERT_EQ(lac.cachedResourceGroupCount(), 1); + ASSERT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); + + auto legacy_group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); + ASSERT_NE(legacy_group, nullptr); + EXPECT_TRUE(legacy_group->enable_gac); + EXPECT_EQ(legacy_group->user_ru_per_sec, 2048); + EXPECT_EQ(legacy_group->group_pb.priority(), ResourceGroup::UserHighPriority); + + auto priority = lac.getPriority(keyspace_id, "default"); + ASSERT_TRUE(priority.has_value()); + EXPECT_EQ(lac.cachedResourceGroupCount(), 1); +} +} // namespace DB::tests From 04608f650d771b121665a081a82d2c567aab0339 Mon Sep 17 00:00:00 2001 From: yongman Date: Mon, 29 Jun 2026 21:02:00 +0800 Subject: [PATCH 2/7] fix default keyspace Signed-off-by: yongman --- .../LocalAdmissionController.cpp | 92 ++++++++++++++++--- .../LocalAdmissionController.h | 8 +- .../gtest_local_admission_controller.cpp | 81 +++++++++++++++- 3 files changed, 160 insertions(+), 21 deletions(-) diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp index 8922740d677..ca79f7d2e73 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp @@ -633,7 +633,7 @@ void LocalAdmissionController::doRequestGAC() const auto resp = cluster->pd_client->acquireTokenBuckets(req); LOG_DEBUG(log, "request to GAC done, req: {}. resp: {}", req.ShortDebugString(), resp.ShortDebugString()); - auto handled = handleTokenBucketsResp(resp); + auto handled = handleTokenBucketsResp(resp, req_rg_names); std::vector> not_found; // not_found includes resource group names that appears in gac_req but not found in resp. @@ -673,7 +673,8 @@ void LocalAdmissionController::checkDegradeMode() } std::vector> LocalAdmissionController::handleTokenBucketsResp( - const resource_manager::TokenBucketsResponse & resp) + const resource_manager::TokenBucketsResponse & resp, + const std::vector> & req_rg_names) { if unlikely (resp.has_error()) { @@ -683,6 +684,13 @@ std::vector> LocalAdmissionController::handle std::vector> handled_resource_group_names; handled_resource_group_names.reserve(resp.responses_size()); + + // Older PD may omit keyspace_id in TokenBucketResponse. Keep request order per resource-group name so + // duplicate names across keyspaces can still be matched to the corresponding requests in order. + std::unordered_map> pending_req_keyspaces_by_name; + for (const auto & [req_keyspace_id, req_name] : req_rg_names) + pending_req_keyspaces_by_name[req_name].push_back(req_keyspace_id); + if (resp.responses().empty()) { LOG_ERROR(log, "got empty TokenBuckets resp from GAC, {}", resp.ShortDebugString()); @@ -699,19 +707,61 @@ std::vector> LocalAdmissionController::handle } const auto & name = one_resp.resource_group_name(); - KeyspaceID keyspace_id = NullspaceID; + KeyspaceID resp_keyspace_id = NullspaceID; + if (one_resp.has_keyspace_id()) + resp_keyspace_id = one_resp.keyspace_id().value(); + + std::optional matched_req_keyspace_id; + ResourceGroupPtr resource_group = nullptr; if (one_resp.has_keyspace_id()) - keyspace_id = one_resp.keyspace_id().value(); - auto resource_group = findResourceGroup(keyspace_id, name); + { + resource_group = findResourceGroup(resp_keyspace_id, name); + } + else + { + auto iter = pending_req_keyspaces_by_name.find(name); + if (iter != pending_req_keyspaces_by_name.end()) + { + for (const auto req_keyspace_id : iter->second) + { + resource_group = findResourceGroup(req_keyspace_id, name); + if (resource_group != nullptr) + { + matched_req_keyspace_id = req_keyspace_id; + break; + } + } + } + if (resource_group == nullptr) + resource_group = findResourceGroup(resp_keyspace_id, name); + } if (resource_group == nullptr) { - LOG_ERROR(log, "cannot find resource group: {}(keyspace={})", name, keyspace_id); + LOG_ERROR(log, "cannot find resource group: {}(keyspace={})", name, resp_keyspace_id); continue; } - handled_resource_group_names.emplace_back(keyspace_id, name); + const KeyspaceID handled_keyspace_id = matched_req_keyspace_id.value_or(resp_keyspace_id); + handled_resource_group_names.emplace_back(handled_keyspace_id, name); + if (auto iter = pending_req_keyspaces_by_name.find(name); iter != pending_req_keyspaces_by_name.end()) + { + auto & pending_keyspaces = iter->second; + if (const auto matched_iter + = std::find(pending_keyspaces.begin(), pending_keyspaces.end(), handled_keyspace_id); + matched_iter != pending_keyspaces.end()) + { + pending_keyspaces.erase(matched_iter); + if (pending_keyspaces.empty()) + pending_req_keyspaces_by_name.erase(iter); + } + } - const String err_msg = fmt::format("handle acquire token resp failed: rg: {}(keyspace={})", name, keyspace_id); + const auto metric_keyspace_id = resource_group->getKeyspaceID(); + const String err_msg = fmt::format( + "handle acquire token resp failed: rg: {}(keyspace={}, cached_keyspace={})", + name, + handled_keyspace_id, + metric_keyspace_id); // It's possible for one_resp.granted_r_u_tokens() to be empty // when the acquire_token_req is only for report RU consumption or GAC got error(like nan token). if (one_resp.granted_r_u_tokens().empty()) @@ -785,7 +835,7 @@ std::vector> LocalAdmissionController::handle // https://github.com/tikv/pd/blob/e9757fbe03260775262763c67f62296fcb26b3c2/pkg/mcs/resourcemanager/server/token_buckets.go#L47 int64_t capacity = granted_token_bucket.granted_tokens().settings().burst_limit(); - const auto name_with_keyspace_id = getResourceGroupMetricName(keyspace_id, name); + const auto name_with_keyspace_id = getResourceGroupMetricName(metric_keyspace_id, name); if (added_tokens > 0) GET_RESOURCE_GROUP_METRIC(tiflash_resource_group, type_gac_resp_tokens, name_with_keyspace_id) .Set(added_tokens); @@ -823,11 +873,21 @@ std::vector> LocalAdmissionController::handle void LocalAdmissionController::watchGACLoop(const std::string & etcd_path) { + auto grpc_context = std::make_unique(); + { + std::lock_guard lock(mu); + active_watch_gac_grpc_contexts.insert(grpc_context.get()); + } + SCOPE_EXIT({ + std::lock_guard lock(mu); + active_watch_gac_grpc_contexts.erase(grpc_context.get()); + }); + while (!stopped.load()) { try { - doWatch(etcd_path); + doWatch(etcd_path, grpc_context.get()); } catch (...) { @@ -848,15 +908,16 @@ void LocalAdmissionController::watchGACLoop(const std::string & etcd_path) })) return; - // Create new grpc_context for each reader/writer. - watch_gac_grpc_context = std::make_unique(); + active_watch_gac_grpc_contexts.erase(grpc_context.get()); + grpc_context = std::make_unique(); + active_watch_gac_grpc_contexts.insert(grpc_context.get()); } } } -void LocalAdmissionController::doWatch(const std::string & etcd_path) +void LocalAdmissionController::doWatch(const std::string & etcd_path, grpc::ClientContext * grpc_context) { - auto stream = etcd_client->watch(watch_gac_grpc_context.get()); + auto stream = etcd_client->watch(grpc_context); auto watch_req = setupWatchReq(etcd_path); LOG_DEBUG(log, "watchGAC req: {}", watch_req.ShortDebugString()); const bool write_ok = stream->Write(watch_req); @@ -1039,7 +1100,8 @@ void LocalAdmissionController::stop() // So still need to lock. { std::lock_guard lock(mu); - watch_gac_grpc_context->TryCancel(); + for (auto * grpc_context : active_watch_gac_grpc_contexts) + grpc_context->TryCancel(); cv.notify_all(); } { diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h index bcea1123c60..dd0c4bc3677 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h @@ -371,7 +371,6 @@ class LocalAdmissionController final : private boost::noncopyable bool start_background_threads = true) : cluster(cluster_) , etcd_client(etcd_client_) - , watch_gac_grpc_context(std::make_unique()) { if (start_background_threads) { @@ -648,7 +647,8 @@ class LocalAdmissionController final : private boost::noncopyable } std::vector> handleTokenBucketsResp( - const resource_manager::TokenBucketsResponse & resp); + const resource_manager::TokenBucketsResponse & resp, + const std::vector> & req_rg_names); static void checkGACRespValid(const resource_manager::ResourceGroup & new_group_pb); @@ -673,7 +673,7 @@ class LocalAdmissionController final : private boost::noncopyable void addReservedDefaultResourceGroup(const KeyspaceID & keyspace_id); // watchGACLoop related methods. - void doWatch(const std::string & etcd_path); + void doWatch(const std::string & etcd_path, grpc::ClientContext * grpc_context); static etcdserverpb::WatchRequest setupWatchReq(const std::string & etcd_path); bool handleDeleteEvent(const std::string & etcd_path, const mvccpb::KeyValue & kv, std::string & err_msg); bool handlePutEvent(const std::string & etcd_path, const mvccpb::KeyValue & kv, std::string & err_msg); @@ -725,7 +725,7 @@ class LocalAdmissionController final : private boost::noncopyable std::atomic need_reset_unique_client_id{false}; uint64_t unique_client_id = 0; Etcd::ClientPtr etcd_client = nullptr; - std::unique_ptr watch_gac_grpc_context = nullptr; + std::unordered_set active_watch_gac_grpc_contexts{}; std::vector background_threads; SteadyClock::time_point current_tick = SteadyClock::time_point::min(); diff --git a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp index ecc36806b95..7d1ba8817fd 100644 --- a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp +++ b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp @@ -26,8 +26,11 @@ class TestPDClient final : public pingcap::pd::MockPDClient public: explicit TestPDClient( std::function - get_resource_group_) + get_resource_group_, + std::function + acquire_token_buckets_ = {}) : get_resource_group(std::move(get_resource_group_)) + , acquire_token_buckets(std::move(acquire_token_buckets_)) {} resource_manager::GetResourceGroupResponse getResourceGroup( @@ -36,14 +39,18 @@ class TestPDClient final : public pingcap::pd::MockPDClient return get_resource_group(req); } - resource_manager::TokenBucketsResponse acquireTokenBuckets(const resource_manager::TokenBucketsRequest &) override + resource_manager::TokenBucketsResponse acquireTokenBuckets(const resource_manager::TokenBucketsRequest & req) override { + if (acquire_token_buckets) + return acquire_token_buckets(req); return {}; } private: std::function get_resource_group; + std::function + acquire_token_buckets; }; } // namespace @@ -144,4 +151,74 @@ TEST(LocalAdmissionControllerTest, KeyspaceLookupFallsBackToLegacyCachedGroup) ASSERT_TRUE(priority.has_value()); EXPECT_EQ(lac.cachedResourceGroupCount(), 1); } + +TEST(LocalAdmissionControllerTest, TokenBucketResponseWithoutKeyspaceMatchesKeyspaceRequest) +{ + constexpr KeyspaceID keyspace_id = 13; + pingcap::kv::Cluster cluster; + cluster.pd_client = std::make_shared( + [](const resource_manager::GetResourceGroupRequest & req) { + resource_manager::GetResourceGroupResponse resp; + if (req.has_keyspace_id()) + { + resp.mutable_error()->set_message("resource group default not found"); + return resp; + } + + auto * group = resp.mutable_group(); + group->set_name(req.resource_group_name()); + group->set_mode(resource_manager::GroupMode::RUMode); + group->set_priority(ResourceGroup::UserHighPriority); + auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); + settings->set_fill_rate(2048); + settings->set_burst_limit(2048); + return resp; + }, + [](const resource_manager::TokenBucketsRequest & req) { + resource_manager::TokenBucketsResponse resp; + EXPECT_EQ(req.requests_size(), 1); + if (req.requests_size() != 1) + return resp; + EXPECT_EQ(req.requests(0).resource_group_name(), "default"); + EXPECT_TRUE(req.requests(0).has_keyspace_id()); + if (!req.requests(0).has_keyspace_id()) + return resp; + EXPECT_EQ(req.requests(0).keyspace_id().value(), keyspace_id); + + auto * one_resp = resp.add_responses(); + one_resp->set_resource_group_name("default"); + auto * granted = one_resp->add_granted_r_u_tokens(); + granted->set_type(resource_manager::RequestUnitType::RU); + granted->set_trickle_time_ms(0); + granted->mutable_granted_tokens()->set_tokens(512); + granted->mutable_granted_tokens()->mutable_settings()->set_burst_limit(2048); + return resp; + }); + + LocalAdmissionController lac(&cluster, nullptr, false, false); + ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); + + auto group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); + ASSERT_NE(group, nullptr); + ASSERT_TRUE(group->enable_gac); + EXPECT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); + + resource_manager::TokenBucketsRequest req; + auto * group_request = req.add_requests(); + group_request->set_resource_group_name("default"); + group_request->mutable_keyspace_id()->set_value(keyspace_id); + auto * ru_items = group_request->mutable_ru_items(); + auto * request_ru = ru_items->add_request_r_u(); + request_ru->set_type(resource_manager::RequestUnitType::RU); + request_ru->set_value(512); + + const auto req_rg_names = std::vector>{{keyspace_id, "default"}}; + auto handled = lac.handleTokenBucketsResp(cluster.pd_client->acquireTokenBuckets(req), req_rg_names); + ASSERT_EQ(handled.size(), 1); + EXPECT_EQ(handled[0], std::make_pair(keyspace_id, std::string("default"))); + + auto priority = lac.getPriority(keyspace_id, "default"); + ASSERT_TRUE(priority.has_value()); + EXPECT_EQ(lac.cachedResourceGroupCount(), 1); +} } // namespace DB::tests From ad9815b5b71963689ca984bbb92a0b77e67fb5c8 Mon Sep 17 00:00:00 2001 From: yongman Date: Wed, 1 Jul 2026 13:21:59 +0800 Subject: [PATCH 3/7] simplify the legacy logic Signed-off-by: yongman --- .../LocalAdmissionController.cpp | 101 +++++-------- .../LocalAdmissionController.h | 3 +- .../gtest_local_admission_controller.cpp | 138 ++++++++++-------- 3 files changed, 113 insertions(+), 129 deletions(-) diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp index ca79f7d2e73..685fa303d8d 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp @@ -19,6 +19,23 @@ namespace DB { +namespace +{ +resource_manager::ResourceGroup normalizeResourceGroupKeyspace( + const resource_manager::ResourceGroup & group_pb, + KeyspaceID keyspace_id) +{ + auto normalized = group_pb; + normalized.mutable_keyspace_id()->set_value(keyspace_id); + return normalized; +} + +KeyspaceID resolveResourceGroupKeyspace(const resource_manager::ResourceGroup & group_pb) +{ + return group_pb.has_keyspace_id() ? group_pb.keyspace_id().value() : NullspaceID; +} +} // namespace + void ResourceGroup::initStaticTokenBucket(int64_t capacity) { std::lock_guard lock(mu); @@ -310,7 +327,8 @@ resource_manager::ResourceGroup LocalAdmissionController::buildReservedDefaultRe void LocalAdmissionController::addReservedDefaultResourceGroup(const KeyspaceID & keyspace_id) { - addResourceGroup(keyspace_id, buildReservedDefaultResourceGroup(keyspace_id), false); + static_cast(keyspace_id); + addResourceGroup(NullspaceID, buildReservedDefaultResourceGroup(NullspaceID), false); } void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & keyspace_id, const std::string & name) @@ -367,9 +385,10 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k if (!legacy_resp.has_error()) { checkGACRespValid(legacy_resp.group()); - const auto rg_keyspace_id - = legacy_resp.group().has_keyspace_id() ? legacy_resp.group().keyspace_id().value() : NullspaceID; - addResourceGroup(rg_keyspace_id, legacy_resp.group()); + const auto resolved_keyspace_id = resolveResourceGroupKeyspace(legacy_resp.group()); + addResourceGroup( + resolved_keyspace_id, + normalizeResourceGroupKeyspace(legacy_resp.group(), resolved_keyspace_id)); return; } if (isReservedDefaultResourceGroup(name)) @@ -396,8 +415,8 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k } checkGACRespValid(resp.group()); - - addResourceGroup(keyspace_id, resp.group()); + const auto resolved_keyspace_id = resolveResourceGroupKeyspace(resp.group()); + addResourceGroup(resolved_keyspace_id, normalizeResourceGroupKeyspace(resp.group(), resolved_keyspace_id)); } void LocalAdmissionController::mainLoop() @@ -633,7 +652,7 @@ void LocalAdmissionController::doRequestGAC() const auto resp = cluster->pd_client->acquireTokenBuckets(req); LOG_DEBUG(log, "request to GAC done, req: {}. resp: {}", req.ShortDebugString(), resp.ShortDebugString()); - auto handled = handleTokenBucketsResp(resp, req_rg_names); + auto handled = handleTokenBucketsResp(resp); std::vector> not_found; // not_found includes resource group names that appears in gac_req but not found in resp. @@ -673,8 +692,7 @@ void LocalAdmissionController::checkDegradeMode() } std::vector> LocalAdmissionController::handleTokenBucketsResp( - const resource_manager::TokenBucketsResponse & resp, - const std::vector> & req_rg_names) + const resource_manager::TokenBucketsResponse & resp) { if unlikely (resp.has_error()) { @@ -684,13 +702,6 @@ std::vector> LocalAdmissionController::handle std::vector> handled_resource_group_names; handled_resource_group_names.reserve(resp.responses_size()); - - // Older PD may omit keyspace_id in TokenBucketResponse. Keep request order per resource-group name so - // duplicate names across keyspaces can still be matched to the corresponding requests in order. - std::unordered_map> pending_req_keyspaces_by_name; - for (const auto & [req_keyspace_id, req_name] : req_rg_names) - pending_req_keyspaces_by_name[req_name].push_back(req_keyspace_id); - if (resp.responses().empty()) { LOG_ERROR(log, "got empty TokenBuckets resp from GAC, {}", resp.ShortDebugString()); @@ -707,61 +718,19 @@ std::vector> LocalAdmissionController::handle } const auto & name = one_resp.resource_group_name(); - KeyspaceID resp_keyspace_id = NullspaceID; + KeyspaceID keyspace_id = NullspaceID; if (one_resp.has_keyspace_id()) - resp_keyspace_id = one_resp.keyspace_id().value(); - - std::optional matched_req_keyspace_id; - ResourceGroupPtr resource_group = nullptr; - if (one_resp.has_keyspace_id()) - { - resource_group = findResourceGroup(resp_keyspace_id, name); - } - else - { - auto iter = pending_req_keyspaces_by_name.find(name); - if (iter != pending_req_keyspaces_by_name.end()) - { - for (const auto req_keyspace_id : iter->second) - { - resource_group = findResourceGroup(req_keyspace_id, name); - if (resource_group != nullptr) - { - matched_req_keyspace_id = req_keyspace_id; - break; - } - } - } - if (resource_group == nullptr) - resource_group = findResourceGroup(resp_keyspace_id, name); - } + keyspace_id = one_resp.keyspace_id().value(); + auto resource_group = findResourceGroup(keyspace_id, name); if (resource_group == nullptr) { - LOG_ERROR(log, "cannot find resource group: {}(keyspace={})", name, resp_keyspace_id); + LOG_ERROR(log, "cannot find resource group: {}(keyspace={})", name, keyspace_id); continue; } - const KeyspaceID handled_keyspace_id = matched_req_keyspace_id.value_or(resp_keyspace_id); - handled_resource_group_names.emplace_back(handled_keyspace_id, name); - if (auto iter = pending_req_keyspaces_by_name.find(name); iter != pending_req_keyspaces_by_name.end()) - { - auto & pending_keyspaces = iter->second; - if (const auto matched_iter - = std::find(pending_keyspaces.begin(), pending_keyspaces.end(), handled_keyspace_id); - matched_iter != pending_keyspaces.end()) - { - pending_keyspaces.erase(matched_iter); - if (pending_keyspaces.empty()) - pending_req_keyspaces_by_name.erase(iter); - } - } + handled_resource_group_names.emplace_back(resource_group->getKeyspaceID(), name); - const auto metric_keyspace_id = resource_group->getKeyspaceID(); - const String err_msg = fmt::format( - "handle acquire token resp failed: rg: {}(keyspace={}, cached_keyspace={})", - name, - handled_keyspace_id, - metric_keyspace_id); + const String err_msg = fmt::format("handle acquire token resp failed: rg: {}(keyspace={})", name, keyspace_id); // It's possible for one_resp.granted_r_u_tokens() to be empty // when the acquire_token_req is only for report RU consumption or GAC got error(like nan token). if (one_resp.granted_r_u_tokens().empty()) @@ -835,7 +804,7 @@ std::vector> LocalAdmissionController::handle // https://github.com/tikv/pd/blob/e9757fbe03260775262763c67f62296fcb26b3c2/pkg/mcs/resourcemanager/server/token_buckets.go#L47 int64_t capacity = granted_token_bucket.granted_tokens().settings().burst_limit(); - const auto name_with_keyspace_id = getResourceGroupMetricName(metric_keyspace_id, name); + const auto name_with_keyspace_id = getResourceGroupMetricName(resource_group->getKeyspaceID(), name); if (added_tokens > 0) GET_RESOURCE_GROUP_METRIC(tiflash_resource_group, type_gac_resp_tokens, name_with_keyspace_id) .Set(added_tokens); @@ -1012,6 +981,8 @@ bool LocalAdmissionController::handlePutEvent( err_msg = "parse pb from etcd value failed"; return false; } + if (keyspace_id == NullspaceID) + group_pb = normalizeResourceGroupKeyspace(group_pb, NullspaceID); { std::lock_guard lock(mu); auto rg = findResourceGroupWithoutLock(keyspace_id, name); diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h index dd0c4bc3677..df09b0910f6 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h @@ -647,8 +647,7 @@ class LocalAdmissionController final : private boost::noncopyable } std::vector> handleTokenBucketsResp( - const resource_manager::TokenBucketsResponse & resp, - const std::vector> & req_rg_names); + const resource_manager::TokenBucketsResponse & resp); static void checkGACRespValid(const resource_manager::ResourceGroup & new_group_pb); diff --git a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp index 7d1ba8817fd..a99b19a5562 100644 --- a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp +++ b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp @@ -68,21 +68,22 @@ TEST(LocalAdmissionControllerTest, ReservedDefaultGroupFallsBackToLocalCacheWhen ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); ASSERT_EQ(lac.cachedResourceGroupCount(), 1); + ASSERT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); auto priority = lac.getPriority(keyspace_id, "default"); ASSERT_TRUE(priority.has_value()); - auto group = lac.getCachedResourceGroupForTest(keyspace_id, "default"); + auto group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); ASSERT_NE(group, nullptr); EXPECT_FALSE(group->enable_gac); EXPECT_EQ(group->user_ru_per_sec, LocalAdmissionController::RESERVED_DEFAULT_RESOURCE_GROUP_RU_PER_SEC); EXPECT_EQ(group->group_pb.priority(), ResourceGroup::UserMediumPriority); EXPECT_TRUE(group->group_pb.has_keyspace_id()); - EXPECT_EQ(group->group_pb.keyspace_id().value(), keyspace_id); + EXPECT_EQ(group->group_pb.keyspace_id().value(), NullspaceID); EXPECT_FALSE(group->buildRequestInfoIfNecessary(SteadyClock::now()).has_value()); } -TEST(LocalAdmissionControllerTest, ReservedDefaultGroupUsesPDConfigWhenPresent) +TEST(LocalAdmissionControllerTest, KeyspaceGroupUsesPDConfigWhenPresent) { constexpr KeyspaceID keyspace_id = 9; pingcap::kv::Cluster cluster; @@ -113,6 +114,37 @@ TEST(LocalAdmissionControllerTest, ReservedDefaultGroupUsesPDConfigWhenPresent) ASSERT_TRUE(priority.has_value()); } +TEST(LocalAdmissionControllerTest, KeyspaceRequestWithoutKeyspaceInPDResponseFallsBackToNullspaceGroup) +{ + constexpr KeyspaceID keyspace_id = 9; + pingcap::kv::Cluster cluster; + cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { + resource_manager::GetResourceGroupResponse resp; + auto * group = resp.mutable_group(); + group->set_name(req.resource_group_name()); + group->set_mode(resource_manager::GroupMode::RUMode); + group->set_priority(ResourceGroup::UserLowPriority); + auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); + settings->set_fill_rate(1024); + settings->set_burst_limit(1024); + return resp; + }); + + LocalAdmissionController lac(&cluster, nullptr, false, false); + + ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); + ASSERT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); + + auto group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); + ASSERT_NE(group, nullptr); + EXPECT_TRUE(group->enable_gac); + EXPECT_TRUE(group->group_pb.has_keyspace_id()); + EXPECT_EQ(group->group_pb.keyspace_id().value(), NullspaceID); + + auto priority = lac.getPriority(keyspace_id, "default"); + ASSERT_TRUE(priority.has_value()); +} + TEST(LocalAdmissionControllerTest, KeyspaceLookupFallsBackToLegacyCachedGroup) { constexpr KeyspaceID keyspace_id = 11; @@ -150,75 +182,57 @@ TEST(LocalAdmissionControllerTest, KeyspaceLookupFallsBackToLegacyCachedGroup) auto priority = lac.getPriority(keyspace_id, "default"); ASSERT_TRUE(priority.has_value()); EXPECT_EQ(lac.cachedResourceGroupCount(), 1); + + auto request_info = legacy_group->buildRequestInfoIfNecessary(SteadyClock::now()); + ASSERT_TRUE(request_info.has_value()); + EXPECT_EQ(request_info->keyspace_id, NullspaceID); } -TEST(LocalAdmissionControllerTest, TokenBucketResponseWithoutKeyspaceMatchesKeyspaceRequest) +TEST(LocalAdmissionControllerTest, LegacyWatchUpdateNormalizesGroupKeyspaceToNullspace) { - constexpr KeyspaceID keyspace_id = 13; + constexpr KeyspaceID keyspace_id = 17; pingcap::kv::Cluster cluster; - cluster.pd_client = std::make_shared( - [](const resource_manager::GetResourceGroupRequest & req) { - resource_manager::GetResourceGroupResponse resp; - if (req.has_keyspace_id()) - { - resp.mutable_error()->set_message("resource group default not found"); - return resp; - } - - auto * group = resp.mutable_group(); - group->set_name(req.resource_group_name()); - group->set_mode(resource_manager::GroupMode::RUMode); - group->set_priority(ResourceGroup::UserHighPriority); - auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); - settings->set_fill_rate(2048); - settings->set_burst_limit(2048); - return resp; - }, - [](const resource_manager::TokenBucketsRequest & req) { - resource_manager::TokenBucketsResponse resp; - EXPECT_EQ(req.requests_size(), 1); - if (req.requests_size() != 1) - return resp; - EXPECT_EQ(req.requests(0).resource_group_name(), "default"); - EXPECT_TRUE(req.requests(0).has_keyspace_id()); - if (!req.requests(0).has_keyspace_id()) - return resp; - EXPECT_EQ(req.requests(0).keyspace_id().value(), keyspace_id); - - auto * one_resp = resp.add_responses(); - one_resp->set_resource_group_name("default"); - auto * granted = one_resp->add_granted_r_u_tokens(); - granted->set_type(resource_manager::RequestUnitType::RU); - granted->set_trickle_time_ms(0); - granted->mutable_granted_tokens()->set_tokens(512); - granted->mutable_granted_tokens()->mutable_settings()->set_burst_limit(2048); + cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { + resource_manager::GetResourceGroupResponse resp; + if (req.has_keyspace_id()) + { + resp.mutable_error()->set_message("resource group default not found"); return resp; - }); + } + + auto * group = resp.mutable_group(); + group->set_name(req.resource_group_name()); + group->set_mode(resource_manager::GroupMode::RUMode); + group->set_priority(ResourceGroup::UserMediumPriority); + auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); + settings->set_fill_rate(1024); + settings->set_burst_limit(1024); + return resp; + }); LocalAdmissionController lac(&cluster, nullptr, false, false); ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); + mvccpb::KeyValue kv; + kv.set_key("resource_group/settings/default"); + resource_manager::ResourceGroup updated_group; + updated_group.set_name("default"); + updated_group.set_mode(resource_manager::GroupMode::RUMode); + updated_group.set_priority(ResourceGroup::UserHighPriority); + updated_group.mutable_keyspace_id()->set_value(keyspace_id); + auto * settings = updated_group.mutable_r_u_settings()->mutable_r_u()->mutable_settings(); + settings->set_fill_rate(4096); + settings->set_burst_limit(4096); + kv.set_value(updated_group.SerializeAsString()); + + std::string err_msg; + ASSERT_TRUE(lac.handlePutEvent(LocalAdmissionController::GAC_RESOURCE_GROUP_ETCD_PATH, kv, err_msg)) << err_msg; + auto group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); ASSERT_NE(group, nullptr); - ASSERT_TRUE(group->enable_gac); - EXPECT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); - - resource_manager::TokenBucketsRequest req; - auto * group_request = req.add_requests(); - group_request->set_resource_group_name("default"); - group_request->mutable_keyspace_id()->set_value(keyspace_id); - auto * ru_items = group_request->mutable_ru_items(); - auto * request_ru = ru_items->add_request_r_u(); - request_ru->set_type(resource_manager::RequestUnitType::RU); - request_ru->set_value(512); - - const auto req_rg_names = std::vector>{{keyspace_id, "default"}}; - auto handled = lac.handleTokenBucketsResp(cluster.pd_client->acquireTokenBuckets(req), req_rg_names); - ASSERT_EQ(handled.size(), 1); - EXPECT_EQ(handled[0], std::make_pair(keyspace_id, std::string("default"))); - - auto priority = lac.getPriority(keyspace_id, "default"); - ASSERT_TRUE(priority.has_value()); - EXPECT_EQ(lac.cachedResourceGroupCount(), 1); + EXPECT_EQ(group->user_ru_per_sec, 4096); + EXPECT_TRUE(group->group_pb.has_keyspace_id()); + EXPECT_EQ(group->group_pb.keyspace_id().value(), NullspaceID); } + } // namespace DB::tests From dd09b4ea969ac9a95857733a86fffa4c5547308c Mon Sep 17 00:00:00 2001 From: yongman Date: Wed, 1 Jul 2026 13:23:48 +0800 Subject: [PATCH 4/7] format Signed-off-by: yongman --- dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp | 3 ++- .../tests/gtest_local_admission_controller.cpp | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp index 685fa303d8d..d26da363b4b 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp @@ -312,7 +312,8 @@ LACRUConsumptionDeltaInfo ResourceGroup::updateRUConsumptionDeltaInfoWithoutLock return info; } -resource_manager::ResourceGroup LocalAdmissionController::buildReservedDefaultResourceGroup(const KeyspaceID & keyspace_id) +resource_manager::ResourceGroup LocalAdmissionController::buildReservedDefaultResourceGroup( + const KeyspaceID & keyspace_id) { resource_manager::ResourceGroup group_pb; group_pb.set_name("default"); diff --git a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp index a99b19a5562..5650075a07a 100644 --- a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp +++ b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp @@ -28,7 +28,8 @@ class TestPDClient final : public pingcap::pd::MockPDClient std::function get_resource_group_, std::function - acquire_token_buckets_ = {}) + acquire_token_buckets_ + = {}) : get_resource_group(std::move(get_resource_group_)) , acquire_token_buckets(std::move(acquire_token_buckets_)) {} @@ -39,7 +40,8 @@ class TestPDClient final : public pingcap::pd::MockPDClient return get_resource_group(req); } - resource_manager::TokenBucketsResponse acquireTokenBuckets(const resource_manager::TokenBucketsRequest & req) override + resource_manager::TokenBucketsResponse acquireTokenBuckets( + const resource_manager::TokenBucketsRequest & req) override { if (acquire_token_buckets) return acquire_token_buckets(req); From 82147c72722d5d4e259479e49ba926ca915af59c Mon Sep 17 00:00:00 2001 From: yongman Date: Wed, 1 Jul 2026 14:46:09 +0800 Subject: [PATCH 5/7] polish Signed-off-by: yongman --- .../LocalAdmissionController.cpp | 102 ++++++++-- .../LocalAdmissionController.h | 10 +- .../gtest_local_admission_controller.cpp | 188 +++++------------- 3 files changed, 143 insertions(+), 157 deletions(-) diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp index d26da363b4b..6c9fbfb1f9e 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp @@ -15,7 +15,10 @@ #include #include +#include +#include #include +#include namespace DB { @@ -34,6 +37,20 @@ KeyspaceID resolveResourceGroupKeyspace(const resource_manager::ResourceGroup & { return group_pb.has_keyspace_id() ? group_pb.keyspace_id().value() : NullspaceID; } + +bool containsIgnoreCase(std::string_view haystack, std::string_view needle) +{ + return std::search( + haystack.begin(), + haystack.end(), + needle.begin(), + needle.end(), + [](char lhs, char rhs) { + return std::tolower(static_cast(lhs)) + == std::tolower(static_cast(rhs)); + }) + != haystack.end(); +} } // namespace void ResourceGroup::initStaticTokenBucket(int64_t capacity) @@ -326,12 +343,21 @@ resource_manager::ResourceGroup LocalAdmissionController::buildReservedDefaultRe return group_pb; } -void LocalAdmissionController::addReservedDefaultResourceGroup(const KeyspaceID & keyspace_id) +void LocalAdmissionController::addReservedDefaultResourceGroup() { - static_cast(keyspace_id); addResourceGroup(NullspaceID, buildReservedDefaultResourceGroup(NullspaceID), false); } +bool LocalAdmissionController::shouldFallbackToLegacyLookup(const resource_manager::GetResourceGroupResponse & resp) +{ + if (!resp.has_error()) + return false; + + const auto & err_msg = resp.error().message(); + return containsIgnoreCase(err_msg, "resource group") + && (containsIgnoreCase(err_msg, "not found") || containsIgnoreCase(err_msg, "does not exist")); +} + void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & keyspace_id, const std::string & name) { if (unlikely(stopped)) @@ -342,7 +368,7 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k { std::lock_guard lock(mu); - if (findResourceGroupWithCompatWithoutLock(keyspace_id, name).group != nullptr) + if (hasExactResourceGroupWithoutLock(keyspace_id, name)) return; } @@ -367,23 +393,35 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k if (resp.has_error()) { + if (!shouldFallbackToLegacyLookup(resp)) + { + RUNTIME_CHECK_MSG( + false, + "warmupResourceGroupInfoCache: {}(keyspace={}) failed: {}", + name, + keyspace_id, + resp.error().message()); + } + resource_manager::GetResourceGroupRequest legacy_req; legacy_req.set_resource_group_name(name); resource_manager::GetResourceGroupResponse legacy_resp; + bool legacy_lookup_failed = false; try { legacy_resp = cluster->pd_client->getResourceGroup(legacy_req); } catch (...) { - LOG_WARNING( + legacy_lookup_failed = true; + tryLogCurrentException( log, - "warmupResourceGroupInfoCache failed to fetch legacy rg for keyspace={}, rg={}, err={}", - keyspace_id, - name, - getCurrentExceptionMessage(false)); + fmt::format( + "warmupResourceGroupInfoCache failed to fetch legacy rg for keyspace={}, rg={}", + keyspace_id, + name)); } - if (!legacy_resp.has_error()) + if (!legacy_lookup_failed && !legacy_resp.has_error()) { checkGACRespValid(legacy_resp.group()); const auto resolved_keyspace_id = resolveResourceGroupKeyspace(legacy_resp.group()); @@ -396,7 +434,7 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k { { std::lock_guard lock(mu); - if (findResourceGroupWithoutLock(NullspaceID, name) != nullptr) + if (findResourceGroupWithCompatWithoutLock(keyspace_id, name).group != nullptr) return; } LOG_WARNING( @@ -404,7 +442,7 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k "warmupResourceGroupInfoCache fallback to reserved default rg for keyspace={}, err={}", keyspace_id, resp.error().message()); - addReservedDefaultResourceGroup(keyspace_id); + addReservedDefaultResourceGroup(); return; } RUNTIME_CHECK_MSG( @@ -984,24 +1022,50 @@ bool LocalAdmissionController::handlePutEvent( } if (keyspace_id == NullspaceID) group_pb = normalizeResourceGroupKeyspace(group_pb, NullspaceID); + bool need_refill = false; { std::lock_guard lock(mu); auto rg = findResourceGroupWithoutLock(keyspace_id, name); + bool inserted_new_exact_group = false; if (rg == nullptr) { - // It happens when query of this resource group has not came. - LOG_DEBUG( - log, - "trying to modify resource group config({}), but cannot find its info", - group_pb.ShortDebugString()); - return true; + const auto compat_rg = findResourceGroupWithCompatWithoutLock(keyspace_id, name); + if (keyspace_id != NullspaceID && compat_rg.group != nullptr && compat_rg.is_fallback) + { + const auto new_user_ru_per_sec = group_pb.r_u_settings().r_u().settings().fill_rate(); + if (max_ru_per_sec < new_user_ru_per_sec) + max_ru_per_sec = new_user_ru_per_sec; + + // A keyspace-scoped watch event comes from PD-managed state, so it should always re-enable GAC + // even if the previous compatibility entry was the local reserved default fallback. + rg = std::make_shared(keyspace_id, group_pb, current_tick, true); + keyspace_resource_groups.insert({{keyspace_id, name}, rg}); + inserted_new_exact_group = true; + need_refill = refill_token_callback != nullptr; + } + else + { + // It happens when query of this resource group has not came. + LOG_DEBUG( + log, + "trying to modify resource group config({}), but cannot find its info", + group_pb.ShortDebugString()); + return true; + } } - else + + if (!inserted_new_exact_group) { + const auto old_user_ru_per_sec = rg->user_ru_per_sec; rg->resetResourceGroup(group_pb); - updateMaxRUPerSecAfterDeleteWithoutLock(rg->user_ru_per_sec); + updateMaxRUPerSecAfterDeleteWithoutLock(old_user_ru_per_sec); + if (max_ru_per_sec < rg->user_ru_per_sec) + max_ru_per_sec = rg->user_ru_per_sec; } + if (need_refill) + refill_token_callback(); } + LOG_INFO(log, "modify resource group {}(keyspace={}) to: {}", name, keyspace_id, group_pb.ShortDebugString()); return true; } diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h index df09b0910f6..27e34919c2b 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h @@ -597,6 +597,10 @@ class LocalAdmissionController final : private boost::noncopyable } return {}; } + bool hasExactResourceGroupWithoutLock(const KeyspaceID & keyspace_id, const std::string & name) const + { + return findResourceGroupWithoutLock(keyspace_id, name) != nullptr; + } void addResourceGroup(const KeyspaceID & keyspace_id, const resource_manager::ResourceGroup & new_group_pb) { @@ -619,8 +623,9 @@ class LocalAdmissionController final : private boost::noncopyable if (iter->second->enable_gac || !enable_gac) return; - updateMaxRUPerSecAfterDeleteWithoutLock(iter->second->user_ru_per_sec); + const auto deletedUserRUPerSec = iter->second->user_ru_per_sec; keyspace_resource_groups.erase(iter); + updateMaxRUPerSecAfterDeleteWithoutLock(deletedUserRUPerSec); } LOG_INFO( @@ -668,8 +673,9 @@ class LocalAdmissionController final : private boost::noncopyable // requestGACLoop related methods. void doRequestGAC(); static bool isReservedDefaultResourceGroup(const std::string & name) { return name == "default"; } + static bool shouldFallbackToLegacyLookup(const resource_manager::GetResourceGroupResponse & resp); static resource_manager::ResourceGroup buildReservedDefaultResourceGroup(const KeyspaceID & keyspace_id); - void addReservedDefaultResourceGroup(const KeyspaceID & keyspace_id); + void addReservedDefaultResourceGroup(); // watchGACLoop related methods. void doWatch(const std::string & etcd_path, grpc::ClientContext * grpc_context); diff --git a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp index 5650075a07a..b42c2ea9558 100644 --- a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp +++ b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include namespace DB::tests @@ -26,12 +27,8 @@ class TestPDClient final : public pingcap::pd::MockPDClient public: explicit TestPDClient( std::function - get_resource_group_, - std::function - acquire_token_buckets_ - = {}) + get_resource_group_) : get_resource_group(std::move(get_resource_group_)) - , acquire_token_buckets(std::move(acquire_token_buckets_)) {} resource_manager::GetResourceGroupResponse getResourceGroup( @@ -40,62 +37,33 @@ class TestPDClient final : public pingcap::pd::MockPDClient return get_resource_group(req); } - resource_manager::TokenBucketsResponse acquireTokenBuckets( - const resource_manager::TokenBucketsRequest & req) override + resource_manager::TokenBucketsResponse acquireTokenBuckets(const resource_manager::TokenBucketsRequest &) override { - if (acquire_token_buckets) - return acquire_token_buckets(req); return {}; } private: std::function get_resource_group; - std::function - acquire_token_buckets; }; } // namespace -TEST(LocalAdmissionControllerTest, ReservedDefaultGroupFallsBackToLocalCacheWhenPDDoesNotPersistIt) +TEST(LocalAdmissionControllerTest, KeyspaceWatchUpdatePromotesLegacyFallbackToExactGroup) { - constexpr KeyspaceID keyspace_id = 7; - pingcap::kv::Cluster cluster; - cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest &) { - resource_manager::GetResourceGroupResponse resp; - resp.mutable_error()->set_message("resource group default not found"); - return resp; - }); - - LocalAdmissionController lac(&cluster, nullptr, false, false); - - ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); - ASSERT_EQ(lac.cachedResourceGroupCount(), 1); - ASSERT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); - - auto priority = lac.getPriority(keyspace_id, "default"); - ASSERT_TRUE(priority.has_value()); - - auto group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); - ASSERT_NE(group, nullptr); - EXPECT_FALSE(group->enable_gac); - EXPECT_EQ(group->user_ru_per_sec, LocalAdmissionController::RESERVED_DEFAULT_RESOURCE_GROUP_RU_PER_SEC); - EXPECT_EQ(group->group_pb.priority(), ResourceGroup::UserMediumPriority); - EXPECT_TRUE(group->group_pb.has_keyspace_id()); - EXPECT_EQ(group->group_pb.keyspace_id().value(), NullspaceID); - EXPECT_FALSE(group->buildRequestInfoIfNecessary(SteadyClock::now()).has_value()); -} - -TEST(LocalAdmissionControllerTest, KeyspaceGroupUsesPDConfigWhenPresent) -{ - constexpr KeyspaceID keyspace_id = 9; + constexpr KeyspaceID keyspace_id = 23; pingcap::kv::Cluster cluster; cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { resource_manager::GetResourceGroupResponse resp; + if (req.has_keyspace_id()) + { + resp.mutable_error()->set_message("resource group default not found"); + return resp; + } + auto * group = resp.mutable_group(); group->set_name(req.resource_group_name()); group->set_mode(resource_manager::GroupMode::RUMode); group->set_priority(ResourceGroup::UserLowPriority); - group->mutable_keyspace_id()->set_value(req.keyspace_id().value()); auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); settings->set_fill_rate(1024); settings->set_burst_limit(1024); @@ -103,138 +71,86 @@ TEST(LocalAdmissionControllerTest, KeyspaceGroupUsesPDConfigWhenPresent) }); LocalAdmissionController lac(&cluster, nullptr, false, false); - ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); - auto group = lac.getCachedResourceGroupForTest(keyspace_id, "default"); - ASSERT_NE(group, nullptr); - EXPECT_TRUE(group->enable_gac); - EXPECT_EQ(group->user_ru_per_sec, 1024); - EXPECT_EQ(group->group_pb.priority(), ResourceGroup::UserLowPriority); - - auto priority = lac.getPriority(keyspace_id, "default"); - ASSERT_TRUE(priority.has_value()); -} - -TEST(LocalAdmissionControllerTest, KeyspaceRequestWithoutKeyspaceInPDResponseFallsBackToNullspaceGroup) -{ - constexpr KeyspaceID keyspace_id = 9; - pingcap::kv::Cluster cluster; - cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { - resource_manager::GetResourceGroupResponse resp; - auto * group = resp.mutable_group(); - group->set_name(req.resource_group_name()); - group->set_mode(resource_manager::GroupMode::RUMode); - group->set_priority(ResourceGroup::UserLowPriority); - auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); - settings->set_fill_rate(1024); - settings->set_burst_limit(1024); - return resp; - }); + ASSERT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); + ASSERT_NE(lac.getCachedResourceGroupForTest(NullspaceID, "default"), nullptr); - LocalAdmissionController lac(&cluster, nullptr, false, false); + mvccpb::KeyValue kv; + kv.set_key("resource_group/keyspace/settings/23/default"); + resource_manager::ResourceGroup updated_group; + updated_group.set_name("default"); + updated_group.set_mode(resource_manager::GroupMode::RUMode); + updated_group.set_priority(ResourceGroup::UserHighPriority); + updated_group.mutable_keyspace_id()->set_value(keyspace_id); + auto * settings = updated_group.mutable_r_u_settings()->mutable_r_u()->mutable_settings(); + settings->set_fill_rate(4096); + settings->set_burst_limit(4096); + kv.set_value(updated_group.SerializeAsString()); - ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); - ASSERT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); + std::string err_msg; + ASSERT_TRUE(lac.handlePutEvent(LocalAdmissionController::GAC_KEYSPACE_RESOURCE_GROUP_ETCD_PATH, kv, err_msg)) + << err_msg; - auto group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); - ASSERT_NE(group, nullptr); - EXPECT_TRUE(group->enable_gac); - EXPECT_TRUE(group->group_pb.has_keyspace_id()); - EXPECT_EQ(group->group_pb.keyspace_id().value(), NullspaceID); + auto exact_group = lac.getCachedResourceGroupForTest(keyspace_id, "default"); + ASSERT_NE(exact_group, nullptr); + EXPECT_TRUE(exact_group->enable_gac); + EXPECT_EQ(exact_group->user_ru_per_sec, 4096); + EXPECT_EQ(exact_group->group_pb.priority(), ResourceGroup::UserHighPriority); - auto priority = lac.getPriority(keyspace_id, "default"); - ASSERT_TRUE(priority.has_value()); + auto request_info = exact_group->buildRequestInfoIfNecessary(SteadyClock::now()); + ASSERT_TRUE(request_info.has_value()); + EXPECT_EQ(request_info->keyspace_id, keyspace_id); } -TEST(LocalAdmissionControllerTest, KeyspaceLookupFallsBackToLegacyCachedGroup) +TEST(LocalAdmissionControllerTest, KeyspaceLookupDoesNotFallbackForNonNotFoundError) { - constexpr KeyspaceID keyspace_id = 11; + constexpr KeyspaceID keyspace_id = 29; pingcap::kv::Cluster cluster; cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { resource_manager::GetResourceGroupResponse resp; if (req.has_keyspace_id()) { - resp.mutable_error()->set_message("resource group default not found"); + resp.mutable_error()->set_message("keyspace metadata is unavailable"); return resp; } auto * group = resp.mutable_group(); group->set_name(req.resource_group_name()); group->set_mode(resource_manager::GroupMode::RUMode); - group->set_priority(ResourceGroup::UserHighPriority); + group->set_priority(ResourceGroup::UserMediumPriority); auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); - settings->set_fill_rate(2048); - settings->set_burst_limit(2048); + settings->set_fill_rate(1024); + settings->set_burst_limit(1024); return resp; }); LocalAdmissionController lac(&cluster, nullptr, false, false); - ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); - ASSERT_EQ(lac.cachedResourceGroupCount(), 1); - ASSERT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); - - auto legacy_group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); - ASSERT_NE(legacy_group, nullptr); - EXPECT_TRUE(legacy_group->enable_gac); - EXPECT_EQ(legacy_group->user_ru_per_sec, 2048); - EXPECT_EQ(legacy_group->group_pb.priority(), ResourceGroup::UserHighPriority); - - auto priority = lac.getPriority(keyspace_id, "default"); - ASSERT_TRUE(priority.has_value()); - EXPECT_EQ(lac.cachedResourceGroupCount(), 1); - - auto request_info = legacy_group->buildRequestInfoIfNecessary(SteadyClock::now()); - ASSERT_TRUE(request_info.has_value()); - EXPECT_EQ(request_info->keyspace_id, NullspaceID); + EXPECT_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default"), DB::Exception); + EXPECT_EQ(lac.cachedResourceGroupCount(), 0); } -TEST(LocalAdmissionControllerTest, LegacyWatchUpdateNormalizesGroupKeyspaceToNullspace) +TEST(LocalAdmissionControllerTest, ReservedDefaultGroupStillFallbacksWhenLegacyLookupThrows) { - constexpr KeyspaceID keyspace_id = 17; + constexpr KeyspaceID keyspace_id = 31; pingcap::kv::Cluster cluster; cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { - resource_manager::GetResourceGroupResponse resp; - if (req.has_keyspace_id()) - { - resp.mutable_error()->set_message("resource group default not found"); - return resp; - } + if (!req.has_keyspace_id()) + throw pingcap::Exception("legacy lookup failed", pingcap::ErrorCodes::UnknownError); - auto * group = resp.mutable_group(); - group->set_name(req.resource_group_name()); - group->set_mode(resource_manager::GroupMode::RUMode); - group->set_priority(ResourceGroup::UserMediumPriority); - auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); - settings->set_fill_rate(1024); - settings->set_burst_limit(1024); + resource_manager::GetResourceGroupResponse resp; + resp.mutable_error()->set_message("resource group default not found"); return resp; }); LocalAdmissionController lac(&cluster, nullptr, false, false); - ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); - - mvccpb::KeyValue kv; - kv.set_key("resource_group/settings/default"); - resource_manager::ResourceGroup updated_group; - updated_group.set_name("default"); - updated_group.set_mode(resource_manager::GroupMode::RUMode); - updated_group.set_priority(ResourceGroup::UserHighPriority); - updated_group.mutable_keyspace_id()->set_value(keyspace_id); - auto * settings = updated_group.mutable_r_u_settings()->mutable_r_u()->mutable_settings(); - settings->set_fill_rate(4096); - settings->set_burst_limit(4096); - kv.set_value(updated_group.SerializeAsString()); - - std::string err_msg; - ASSERT_TRUE(lac.handlePutEvent(LocalAdmissionController::GAC_RESOURCE_GROUP_ETCD_PATH, kv, err_msg)) << err_msg; + ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); auto group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); ASSERT_NE(group, nullptr); - EXPECT_EQ(group->user_ru_per_sec, 4096); - EXPECT_TRUE(group->group_pb.has_keyspace_id()); - EXPECT_EQ(group->group_pb.keyspace_id().value(), NullspaceID); + EXPECT_FALSE(group->enable_gac); + EXPECT_EQ(group->user_ru_per_sec, LocalAdmissionController::RESERVED_DEFAULT_RESOURCE_GROUP_RU_PER_SEC); } } // namespace DB::tests From add89e110a5037b5b92e78694851800cf230c0a1 Mon Sep 17 00:00:00 2001 From: yongman Date: Fri, 3 Jul 2026 21:49:20 +0800 Subject: [PATCH 6/7] refactor for simplify Signed-off-by: yongman --- .../LocalAdmissionController.cpp | 262 +++++++----------- .../LocalAdmissionController.h | 97 +++---- .../gtest_local_admission_controller.cpp | 160 +++++------ 3 files changed, 220 insertions(+), 299 deletions(-) diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp index 6c9fbfb1f9e..cc26aaee534 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp @@ -18,39 +18,15 @@ #include #include #include -#include namespace DB { namespace { -resource_manager::ResourceGroup normalizeResourceGroupKeyspace( - const resource_manager::ResourceGroup & group_pb, - KeyspaceID keyspace_id) -{ - auto normalized = group_pb; - normalized.mutable_keyspace_id()->set_value(keyspace_id); - return normalized; -} - KeyspaceID resolveResourceGroupKeyspace(const resource_manager::ResourceGroup & group_pb) { return group_pb.has_keyspace_id() ? group_pb.keyspace_id().value() : NullspaceID; } - -bool containsIgnoreCase(std::string_view haystack, std::string_view needle) -{ - return std::search( - haystack.begin(), - haystack.end(), - needle.begin(), - needle.end(), - [](char lhs, char rhs) { - return std::tolower(static_cast(lhs)) - == std::tolower(static_cast(rhs)); - }) - != haystack.end(); -} } // namespace void ResourceGroup::initStaticTokenBucket(int64_t capacity) @@ -93,9 +69,6 @@ uint64_t ResourceGroup::getPriority(uint64_t max_ru_per_sec) const std::optional ResourceGroup::buildRequestInfoIfNecessary(const SteadyClock::time_point & now) { - if (!enable_gac) - return {}; - std::lock_guard lock(mu); if (!beginRequestWithoutLock(now)) { @@ -329,35 +302,6 @@ LACRUConsumptionDeltaInfo ResourceGroup::updateRUConsumptionDeltaInfoWithoutLock return info; } -resource_manager::ResourceGroup LocalAdmissionController::buildReservedDefaultResourceGroup( - const KeyspaceID & keyspace_id) -{ - resource_manager::ResourceGroup group_pb; - group_pb.set_name("default"); - group_pb.set_mode(resource_manager::GroupMode::RUMode); - group_pb.set_priority(ResourceGroup::UserMediumPriority); - group_pb.mutable_keyspace_id()->set_value(keyspace_id); - auto * settings = group_pb.mutable_r_u_settings()->mutable_r_u()->mutable_settings(); - settings->set_fill_rate(RESERVED_DEFAULT_RESOURCE_GROUP_RU_PER_SEC); - settings->set_burst_limit(-1); - return group_pb; -} - -void LocalAdmissionController::addReservedDefaultResourceGroup() -{ - addResourceGroup(NullspaceID, buildReservedDefaultResourceGroup(NullspaceID), false); -} - -bool LocalAdmissionController::shouldFallbackToLegacyLookup(const resource_manager::GetResourceGroupResponse & resp) -{ - if (!resp.has_error()) - return false; - - const auto & err_msg = resp.error().message(); - return containsIgnoreCase(err_msg, "resource group") - && (containsIgnoreCase(err_msg, "not found") || containsIgnoreCase(err_msg, "does not exist")); -} - void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & keyspace_id, const std::string & name) { if (unlikely(stopped)) @@ -368,12 +312,19 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k { std::lock_guard lock(mu); - if (hasExactResourceGroupWithoutLock(keyspace_id, name)) + if (findResourceGroupWithCompatWithoutLock(keyspace_id, name) != nullptr) return; } resource_manager::GetResourceGroupRequest req; - req.mutable_keyspace_id()->set_value(keyspace_id); + { + std::lock_guard lock(mu); + // Until the backend is identified, try the keyspace-aware request shape first. pd-cse + // ignores the field and still returns the shared legacy definition, while new PD returns + // a keyspace-scoped group with keyspace_id set. + if (with_keyspace && resource_group_backend_mode != ResourceGroupBackendMode::LegacyGlobal) + req.mutable_keyspace_id()->set_value(keyspace_id); + } req.set_resource_group_name(name); resource_manager::GetResourceGroupResponse resp; @@ -392,70 +343,23 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k } if (resp.has_error()) - { - if (!shouldFallbackToLegacyLookup(resp)) - { - RUNTIME_CHECK_MSG( - false, - "warmupResourceGroupInfoCache: {}(keyspace={}) failed: {}", - name, - keyspace_id, - resp.error().message()); - } - - resource_manager::GetResourceGroupRequest legacy_req; - legacy_req.set_resource_group_name(name); - resource_manager::GetResourceGroupResponse legacy_resp; - bool legacy_lookup_failed = false; - try - { - legacy_resp = cluster->pd_client->getResourceGroup(legacy_req); - } - catch (...) - { - legacy_lookup_failed = true; - tryLogCurrentException( - log, - fmt::format( - "warmupResourceGroupInfoCache failed to fetch legacy rg for keyspace={}, rg={}", - keyspace_id, - name)); - } - if (!legacy_lookup_failed && !legacy_resp.has_error()) - { - checkGACRespValid(legacy_resp.group()); - const auto resolved_keyspace_id = resolveResourceGroupKeyspace(legacy_resp.group()); - addResourceGroup( - resolved_keyspace_id, - normalizeResourceGroupKeyspace(legacy_resp.group(), resolved_keyspace_id)); - return; - } - if (isReservedDefaultResourceGroup(name)) - { - { - std::lock_guard lock(mu); - if (findResourceGroupWithCompatWithoutLock(keyspace_id, name).group != nullptr) - return; - } - LOG_WARNING( - log, - "warmupResourceGroupInfoCache fallback to reserved default rg for keyspace={}, err={}", - keyspace_id, - resp.error().message()); - addReservedDefaultResourceGroup(); - return; - } RUNTIME_CHECK_MSG( false, "warmupResourceGroupInfoCache: {}(keyspace={}) failed: {}", name, keyspace_id, resp.error().message()); - } checkGACRespValid(resp.group()); - const auto resolved_keyspace_id = resolveResourceGroupKeyspace(resp.group()); - addResourceGroup(resolved_keyspace_id, normalizeResourceGroupKeyspace(resp.group(), resolved_keyspace_id)); + const auto detected_mode + = with_keyspace ? detectBackendMode(resp.group()) : ResourceGroupBackendMode::LegacyGlobal; + // pd-cse returns a group without keyspace_id and keeps using the shared legacy namespace. + // New PD returns the concrete keyspace-scoped definition. TiFlash only needs this one bit to + // choose its cache key shape and which watch prefix to keep following afterwards. + updateBackendMode(detected_mode); + const auto resolved_keyspace_id + = detected_mode == ResourceGroupBackendMode::LegacyGlobal ? NullspaceID : resolveResourceGroupKeyspace(resp.group()); + addResourceGroup(resolved_keyspace_id, resp.group()); } void LocalAdmissionController::mainLoop() @@ -593,6 +497,8 @@ std::optional LocalAdmissionController::b for (const auto & info : request_infos) { auto * group_request = gac_req.add_requests(); + // Nullspace/legacy requests were encoded without keyspace_id. Preserve that wire format so + // the request can still be understood by older PD/GAC deployments. if (info.keyspace_id != NullspaceID) group_request->mutable_keyspace_id()->set_value(info.keyspace_id); group_request->set_resource_group_name(info.resource_group_name); @@ -659,6 +565,7 @@ static std::vector> extractGACReqNames( std::vector> res; res.reserve(gac_req.requests_size()); for (const auto & req : gac_req.requests()) + // Must mirror buildGACRequest(): a missing field means the logical key is NullspaceID. res.push_back({req.has_keyspace_id() ? req.keyspace_id().value() : NullspaceID, req.resource_group_name()}); return res; } @@ -884,11 +791,12 @@ void LocalAdmissionController::watchGACLoop(const std::string & etcd_path) auto grpc_context = std::make_unique(); { std::lock_guard lock(mu); - active_watch_gac_grpc_contexts.insert(grpc_context.get()); + active_watch_gac_grpc_context = grpc_context.get(); } SCOPE_EXIT({ std::lock_guard lock(mu); - active_watch_gac_grpc_contexts.erase(grpc_context.get()); + if (active_watch_gac_grpc_context == grpc_context.get()) + active_watch_gac_grpc_context = nullptr; }); while (!stopped.load()) @@ -916,9 +824,10 @@ void LocalAdmissionController::watchGACLoop(const std::string & etcd_path) })) return; - active_watch_gac_grpc_contexts.erase(grpc_context.get()); + if (active_watch_gac_grpc_context == grpc_context.get()) + active_watch_gac_grpc_context = nullptr; grpc_context = std::make_unique(); - active_watch_gac_grpc_contexts.insert(grpc_context.get()); + active_watch_gac_grpc_context = grpc_context.get(); } } } @@ -1020,49 +929,21 @@ bool LocalAdmissionController::handlePutEvent( err_msg = "parse pb from etcd value failed"; return false; } - if (keyspace_id == NullspaceID) - group_pb = normalizeResourceGroupKeyspace(group_pb, NullspaceID); - bool need_refill = false; { std::lock_guard lock(mu); auto rg = findResourceGroupWithoutLock(keyspace_id, name); - bool inserted_new_exact_group = false; if (rg == nullptr) { - const auto compat_rg = findResourceGroupWithCompatWithoutLock(keyspace_id, name); - if (keyspace_id != NullspaceID && compat_rg.group != nullptr && compat_rg.is_fallback) - { - const auto new_user_ru_per_sec = group_pb.r_u_settings().r_u().settings().fill_rate(); - if (max_ru_per_sec < new_user_ru_per_sec) - max_ru_per_sec = new_user_ru_per_sec; - - // A keyspace-scoped watch event comes from PD-managed state, so it should always re-enable GAC - // even if the previous compatibility entry was the local reserved default fallback. - rg = std::make_shared(keyspace_id, group_pb, current_tick, true); - keyspace_resource_groups.insert({{keyspace_id, name}, rg}); - inserted_new_exact_group = true; - need_refill = refill_token_callback != nullptr; - } - else - { - // It happens when query of this resource group has not came. - LOG_DEBUG( - log, - "trying to modify resource group config({}), but cannot find its info", - group_pb.ShortDebugString()); - return true; - } - } - - if (!inserted_new_exact_group) - { - const auto old_user_ru_per_sec = rg->user_ru_per_sec; - rg->resetResourceGroup(group_pb); - updateMaxRUPerSecAfterDeleteWithoutLock(old_user_ru_per_sec); - if (max_ru_per_sec < rg->user_ru_per_sec) - max_ru_per_sec = rg->user_ru_per_sec; + // Ignore updates for groups that have never been warmed into the local cache yet. + LOG_DEBUG(log, "trying to modify resource group config({}), but cannot find its info", group_pb.ShortDebugString()); + return true; } - if (need_refill) + const auto old_user_ru_per_sec = rg->user_ru_per_sec; + rg->resetResourceGroup(group_pb); + updateMaxRUPerSecAfterDeleteWithoutLock(old_user_ru_per_sec); + if (max_ru_per_sec < rg->user_ru_per_sec) + max_ru_per_sec = rg->user_ru_per_sec; + if (refill_token_callback) refill_token_callback(); } @@ -1077,8 +958,9 @@ bool LocalAdmissionController::parseResourceGroupNameFromWatchKey( std::string & parsed_rg_name, std::string & err_msg) { - // Expect etcd_key: resource_group/settings/rg_name OR resource_group/settings/keyspace_id/rg_name - // etcd_key_prefix is resource_group/settings + // Expect: + // 1. resource_group/settings/ + // 2. resource_group/keyspace/settings// if (etcd_key.length() <= etcd_key_prefix.length() + 1) { err_msg = fmt::format("expect etcd key: {}/resource_group_name, but got {}", etcd_key_prefix, etcd_key); @@ -1136,15 +1018,20 @@ void LocalAdmissionController::stop() // So still need to lock. { std::lock_guard lock(mu); - for (auto * grpc_context : active_watch_gac_grpc_contexts) - grpc_context->TryCancel(); + if (active_watch_gac_grpc_context != nullptr) + active_watch_gac_grpc_context->TryCancel(); cv.notify_all(); } { std::lock_guard lock(gac_requests_mu); gac_requests_cv.notify_all(); } - for (auto & thread : background_threads) + std::vector threads_to_join; + { + std::lock_guard lock(mu); + threads_to_join = std::move(background_threads); + } + for (auto & thread : threads_to_join) { if (thread.joinable()) thread.join(); @@ -1191,6 +1078,59 @@ std::unique_ptr LocalAdmissionController::global_insta const std::string LocalAdmissionController::GAC_RESOURCE_GROUP_ETCD_PATH = "resource_group/settings"; const std::string LocalAdmissionController::GAC_KEYSPACE_RESOURCE_GROUP_ETCD_PATH = "resource_group/keyspace/settings"; +ResourceGroupBackendMode LocalAdmissionController::detectBackendMode(const resource_manager::ResourceGroup & group_pb) +{ + return group_pb.has_keyspace_id() ? ResourceGroupBackendMode::KeyspaceScoped + : ResourceGroupBackendMode::LegacyGlobal; +} + +const std::string & LocalAdmissionController::getWatchEtcdPath(ResourceGroupBackendMode mode) +{ + switch (mode) + { + case ResourceGroupBackendMode::LegacyGlobal: + return GAC_RESOURCE_GROUP_ETCD_PATH; + case ResourceGroupBackendMode::KeyspaceScoped: + return GAC_KEYSPACE_RESOURCE_GROUP_ETCD_PATH; + case ResourceGroupBackendMode::Unknown: + throw Exception(ErrorCodes::LOGICAL_ERROR, "watch path is unavailable for unknown backend mode"); + } + + throw Exception(ErrorCodes::LOGICAL_ERROR, "unexpected backend mode {}", magic_enum::enum_name(mode)); +} + +void LocalAdmissionController::updateBackendMode(ResourceGroupBackendMode detected_mode) +{ + std::lock_guard lock(mu); + if (detected_mode == ResourceGroupBackendMode::Unknown) + return; + + if (resource_group_backend_mode == detected_mode) + return; + + if (resource_group_backend_mode != ResourceGroupBackendMode::Unknown) + { + LOG_WARNING( + log, + "keep resource-group backend mode {} and ignore newly detected mode {}", + magic_enum::enum_name(resource_group_backend_mode), + magic_enum::enum_name(detected_mode)); + return; + } + + resource_group_backend_mode = detected_mode; + LOG_INFO(log, "resource-group backend mode resolved to {}", magic_enum::enum_name(resource_group_backend_mode)); + + if (start_background_threads && !watch_thread_started && !stopped.load()) + { + // Only keep one watch path alive. After the backend mode is known there is no need to + // continue listening on both legacy and keyspace prefixes at the same time. + const auto etcd_path = getWatchEtcdPath(resource_group_backend_mode); + background_threads.emplace_back([this, etcd_path] { this->watchGACLoop(etcd_path); }); + watch_thread_started = true; + } +} + bool LocalAdmissionController::parseKeyspaceEtcdKey( const std::string & etcd_key_prefix, const std::string & etcd_key, @@ -1198,7 +1138,7 @@ bool LocalAdmissionController::parseKeyspaceEtcdKey( std::string & parsed_rg_name, std::string & err_msg) { - // resource_group/settings/keyspace_id/rg_name -> keyspace_id/rg_name + // resource_group/keyspace/settings// -> / auto tmp_str = std::string(etcd_key.begin() + etcd_key_prefix.length() + 1, etcd_key.end()); size_t slash_pos = tmp_str.find('/'); if (slash_pos == std::string::npos) diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h index 27e34919c2b..faa5d1b71b0 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h @@ -46,12 +46,6 @@ struct GACRequestInfo double ru_consumption_delta; }; -struct ResourceGroupCompatLookupResult -{ - std::shared_ptr group; - bool is_fallback = false; -}; - inline std::string getResourceGroupMetricName(const KeyspaceID & keyspace_id, const std::string & name) { return fmt::format("{}-{}", keyspace_id, name); @@ -80,12 +74,10 @@ class ResourceGroup final : private boost::noncopyable explicit ResourceGroup( const KeyspaceID & keyspace_id_, const resource_manager::ResourceGroup & group_pb_, - const SteadyClock::time_point & tp, - bool enable_gac_ = true) + const SteadyClock::time_point & tp) : keyspace_id(keyspace_id_) , name(group_pb_.name()) , name_with_keyspace_id(getResourceGroupMetricName(keyspace_id_, group_pb_.name())) - , enable_gac(enable_gac_) , group_pb(group_pb_) , log(Logger::get("resource group:" + group_pb_.name())) { @@ -307,7 +299,6 @@ class ResourceGroup final : private boost::noncopyable const KeyspaceID keyspace_id; const std::string name; const std::string name_with_keyspace_id; // For metrics. - bool enable_gac = true; uint32_t user_priority_val = 0; uint64_t user_ru_per_sec = 0; bool burstable = false; @@ -356,6 +347,13 @@ struct LACPairHash } }; +enum class ResourceGroupBackendMode +{ + Unknown, + LegacyGlobal, + KeyspaceScoped, +}; + // LocalAdmissionController is the local(tiflash) part of the distributed token bucket algorithm. // It manages all resource groups: // 1. Creation, deletion and config updates of resource group. @@ -371,24 +369,21 @@ class LocalAdmissionController final : private boost::noncopyable bool start_background_threads = true) : cluster(cluster_) , etcd_client(etcd_client_) + , with_keyspace(with_keyspace) + , start_background_threads(start_background_threads) + , resource_group_backend_mode(with_keyspace ? ResourceGroupBackendMode::Unknown : ResourceGroupBackendMode::LegacyGlobal) { + current_tick = SteadyClock::now(); + last_clear_cpu_time = current_tick; + if (start_background_threads) { background_threads.emplace_back([this] { this->mainLoop(); }); background_threads.emplace_back([this] { this->requestGACLoop(); }); - if (with_keyspace) - { - background_threads.emplace_back([this] { this->watchGACLoop(GAC_KEYSPACE_RESOURCE_GROUP_ETCD_PATH); }); + if (resource_group_backend_mode == ResourceGroupBackendMode::LegacyGlobal) background_threads.emplace_back([this] { this->watchGACLoop(GAC_RESOURCE_GROUP_ETCD_PATH); }); - } - else - { - background_threads.emplace_back([this] { this->watchGACLoop(GAC_RESOURCE_GROUP_ETCD_PATH); }); - } + watch_thread_started = (resource_group_backend_mode != ResourceGroupBackendMode::Unknown); } - - current_tick = SteadyClock::now(); - last_clear_cpu_time = current_tick; } ~LocalAdmissionController() { safeStop(); } @@ -510,7 +505,6 @@ class LocalAdmissionController final : private boost::noncopyable static constexpr auto DEGRADE_MODE_DURATION = std::chrono::seconds(120); static constexpr double ACQUIRE_RU_AMPLIFICATION = 1.1; static constexpr auto DEFAULT_MAX_EST_WAIT_DURATION = std::chrono::milliseconds(1000); - static constexpr uint64_t RESERVED_DEFAULT_RESOURCE_GROUP_RU_PER_SEC = std::numeric_limits::max(); #ifndef DBMS_PUBLIC_GTEST private: @@ -566,14 +560,14 @@ class LocalAdmissionController final : private boost::noncopyable ResourceGroupPtr findResourceGroup(const KeyspaceID & keyspace_id, const std::string & name) const { std::lock_guard lock(mu); - return findResourceGroupWithCompatWithoutLock(keyspace_id, name).group; + return findResourceGroupWithCompatWithoutLock(keyspace_id, name); } std::pair findResourceGroupAndMaxRUPerSec( const KeyspaceID & keyspace_id, const std::string & name) const { std::lock_guard lock(mu); - auto rg = findResourceGroupWithCompatWithoutLock(keyspace_id, name).group; + auto rg = findResourceGroupWithCompatWithoutLock(keyspace_id, name); return {rg, max_ru_per_sec}; } ResourceGroupPtr findResourceGroupWithoutLock(const KeyspaceID & keyspace_id, const std::string & name) const @@ -583,33 +577,22 @@ class LocalAdmissionController final : private boost::noncopyable return nullptr; return iter->second; } - ResourceGroupCompatLookupResult findResourceGroupWithCompatWithoutLock( - const KeyspaceID & keyspace_id, - const std::string & name) const + ResourceGroupPtr findResourceGroupWithCompatWithoutLock(const KeyspaceID & keyspace_id, const std::string & name) const { if (auto group = findResourceGroupWithoutLock(keyspace_id, name); likely(group != nullptr)) - return {.group = group, .is_fallback = false}; + return group; - if (keyspace_id != NullspaceID) + // pd-cse only has one shared resource-group namespace, so keyed requests must fall back + // to the Nullspace cache once we know we are talking to that backend. + if (resource_group_backend_mode == ResourceGroupBackendMode::LegacyGlobal && keyspace_id != NullspaceID) { if (auto group = findResourceGroupWithoutLock(NullspaceID, name); likely(group != nullptr)) - return {.group = group, .is_fallback = true}; + return group; } - return {}; - } - bool hasExactResourceGroupWithoutLock(const KeyspaceID & keyspace_id, const std::string & name) const - { - return findResourceGroupWithoutLock(keyspace_id, name) != nullptr; + return nullptr; } void addResourceGroup(const KeyspaceID & keyspace_id, const resource_manager::ResourceGroup & new_group_pb) - { - addResourceGroup(keyspace_id, new_group_pb, true); - } - void addResourceGroup( - const KeyspaceID & keyspace_id, - const resource_manager::ResourceGroup & new_group_pb, - bool enable_gac) { uint64_t user_ru_per_sec = new_group_pb.r_u_settings().r_u().settings().fill_rate(); std::lock_guard lock(mu); @@ -619,22 +602,14 @@ class LocalAdmissionController final : private boost::noncopyable auto iter = keyspace_resource_groups.find({keyspace_id, new_group_pb.name()}); if (iter != keyspace_resource_groups.end()) - { - if (iter->second->enable_gac || !enable_gac) - return; - - const auto deletedUserRUPerSec = iter->second->user_ru_per_sec; - keyspace_resource_groups.erase(iter); - updateMaxRUPerSecAfterDeleteWithoutLock(deletedUserRUPerSec); - } + return; LOG_INFO( log, - "add new resource group, keyspace={} enable_gac={} info: {}", + "add new resource group, keyspace={} info: {}", keyspace_id, - enable_gac, new_group_pb.ShortDebugString()); - auto new_group = std::make_shared(keyspace_id, new_group_pb, current_tick, enable_gac); + auto new_group = std::make_shared(keyspace_id, new_group_pb, current_tick); keyspace_resource_groups.insert({{keyspace_id, new_group_pb.name()}, new_group}); if (refill_token_callback) @@ -672,10 +647,9 @@ class LocalAdmissionController final : private boost::noncopyable // requestGACLoop related methods. void doRequestGAC(); - static bool isReservedDefaultResourceGroup(const std::string & name) { return name == "default"; } - static bool shouldFallbackToLegacyLookup(const resource_manager::GetResourceGroupResponse & resp); - static resource_manager::ResourceGroup buildReservedDefaultResourceGroup(const KeyspaceID & keyspace_id); - void addReservedDefaultResourceGroup(); + void updateBackendMode(ResourceGroupBackendMode detected_mode); + static ResourceGroupBackendMode detectBackendMode(const resource_manager::ResourceGroup & group_pb); + static const std::string & getWatchEtcdPath(ResourceGroupBackendMode mode); // watchGACLoop related methods. void doWatch(const std::string & etcd_path, grpc::ClientContext * grpc_context); @@ -730,7 +704,14 @@ class LocalAdmissionController final : private boost::noncopyable std::atomic need_reset_unique_client_id{false}; uint64_t unique_client_id = 0; Etcd::ClientPtr etcd_client = nullptr; - std::unordered_set active_watch_gac_grpc_contexts{}; + const bool with_keyspace; + const bool start_background_threads; + // Unknown means TiFlash has not seen any successful GetResourceGroup response yet. Once a + // response arrives, its keyspace_id presence decides whether we should operate in pd-cse's + // shared namespace or PD's keyspace-scoped namespace. + ResourceGroupBackendMode resource_group_backend_mode; + grpc::ClientContext * active_watch_gac_grpc_context = nullptr; + bool watch_thread_started = false; std::vector background_threads; SteadyClock::time_point current_tick = SteadyClock::time_point::min(); diff --git a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp index b42c2ea9558..ef13ea89a8d 100644 --- a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp +++ b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp @@ -13,11 +13,12 @@ // limitations under the License. #include -#include #include -#include #include +#include +#include + namespace DB::tests { namespace @@ -46,111 +47,110 @@ class TestPDClient final : public pingcap::pd::MockPDClient std::function get_resource_group; }; + +resource_manager::GetResourceGroupResponse buildResourceGroupResp( + const resource_manager::GetResourceGroupRequest & req, + uint64_t fill_rate, + bool with_keyspace_id) +{ + resource_manager::GetResourceGroupResponse resp; + auto * group = resp.mutable_group(); + group->set_name(req.resource_group_name()); + group->set_mode(resource_manager::GroupMode::RUMode); + group->set_priority(ResourceGroup::UserLowPriority); + if (with_keyspace_id) + group->mutable_keyspace_id()->set_value(req.keyspace_id().value()); + auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); + settings->set_fill_rate(fill_rate); + settings->set_burst_limit(fill_rate); + return resp; +} } // namespace -TEST(LocalAdmissionControllerTest, KeyspaceWatchUpdatePromotesLegacyFallbackToExactGroup) +TEST(LocalAdmissionControllerTest, LegacyBackendCachesSharedGroupAndOmitsKeyspaceAfterDetection) { constexpr KeyspaceID keyspace_id = 23; + std::vector request_has_keyspace; pingcap::kv::Cluster cluster; - cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { - resource_manager::GetResourceGroupResponse resp; - if (req.has_keyspace_id()) - { - resp.mutable_error()->set_message("resource group default not found"); - return resp; - } - - auto * group = resp.mutable_group(); - group->set_name(req.resource_group_name()); - group->set_mode(resource_manager::GroupMode::RUMode); - group->set_priority(ResourceGroup::UserLowPriority); - auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); - settings->set_fill_rate(1024); - settings->set_burst_limit(1024); - return resp; - }); + cluster.pd_client = std::make_shared( + [&request_has_keyspace](const resource_manager::GetResourceGroupRequest & req) { + request_has_keyspace.push_back(req.has_keyspace_id()); + const auto fill_rate = req.resource_group_name() == "default" ? 1024 : 2048; + return buildResourceGroupResp(req, fill_rate, false); + }); + + LocalAdmissionController lac(&cluster, nullptr, true, false); - LocalAdmissionController lac(&cluster, nullptr, false, false); ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); + ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id + 1, "analytics")); + + ASSERT_EQ(request_has_keyspace.size(), 2); + EXPECT_TRUE(request_has_keyspace[0]); + EXPECT_FALSE(request_has_keyspace[1]); - ASSERT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); - ASSERT_NE(lac.getCachedResourceGroupForTest(NullspaceID, "default"), nullptr); - - mvccpb::KeyValue kv; - kv.set_key("resource_group/keyspace/settings/23/default"); - resource_manager::ResourceGroup updated_group; - updated_group.set_name("default"); - updated_group.set_mode(resource_manager::GroupMode::RUMode); - updated_group.set_priority(ResourceGroup::UserHighPriority); - updated_group.mutable_keyspace_id()->set_value(keyspace_id); - auto * settings = updated_group.mutable_r_u_settings()->mutable_r_u()->mutable_settings(); - settings->set_fill_rate(4096); - settings->set_burst_limit(4096); - kv.set_value(updated_group.SerializeAsString()); - - std::string err_msg; - ASSERT_TRUE(lac.handlePutEvent(LocalAdmissionController::GAC_KEYSPACE_RESOURCE_GROUP_ETCD_PATH, kv, err_msg)) - << err_msg; - - auto exact_group = lac.getCachedResourceGroupForTest(keyspace_id, "default"); - ASSERT_NE(exact_group, nullptr); - EXPECT_TRUE(exact_group->enable_gac); - EXPECT_EQ(exact_group->user_ru_per_sec, 4096); - EXPECT_EQ(exact_group->group_pb.priority(), ResourceGroup::UserHighPriority); - - auto request_info = exact_group->buildRequestInfoIfNecessary(SteadyClock::now()); + EXPECT_EQ(lac.getCachedResourceGroupForTest(keyspace_id, "default"), nullptr); + auto default_group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); + ASSERT_NE(default_group, nullptr); + EXPECT_TRUE(lac.getPriority(keyspace_id, "default").has_value()); + + auto request_info = default_group->buildRequestInfoIfNecessary(SteadyClock::now()); ASSERT_TRUE(request_info.has_value()); - EXPECT_EQ(request_info->keyspace_id, keyspace_id); + EXPECT_EQ(request_info->keyspace_id, NullspaceID); + + auto analytics_group = lac.getCachedResourceGroupForTest(NullspaceID, "analytics"); + ASSERT_NE(analytics_group, nullptr); + EXPECT_EQ(analytics_group->user_ru_per_sec, 2048); } -TEST(LocalAdmissionControllerTest, KeyspaceLookupDoesNotFallbackForNonNotFoundError) +TEST(LocalAdmissionControllerTest, KeyspaceScopedBackendCachesExactGroupAndKeepsKeyspaceRequests) { constexpr KeyspaceID keyspace_id = 29; + std::vector request_has_keyspace; pingcap::kv::Cluster cluster; - cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { - resource_manager::GetResourceGroupResponse resp; - if (req.has_keyspace_id()) - { - resp.mutable_error()->set_message("keyspace metadata is unavailable"); - return resp; - } - - auto * group = resp.mutable_group(); - group->set_name(req.resource_group_name()); - group->set_mode(resource_manager::GroupMode::RUMode); - group->set_priority(ResourceGroup::UserMediumPriority); - auto * settings = group->mutable_r_u_settings()->mutable_r_u()->mutable_settings(); - settings->set_fill_rate(1024); - settings->set_burst_limit(1024); - return resp; - }); + cluster.pd_client = std::make_shared( + [&request_has_keyspace](const resource_manager::GetResourceGroupRequest & req) { + request_has_keyspace.push_back(req.has_keyspace_id()); + return buildResourceGroupResp(req, 4096, true); + }); - LocalAdmissionController lac(&cluster, nullptr, false, false); + LocalAdmissionController lac(&cluster, nullptr, true, false); - EXPECT_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default"), DB::Exception); - EXPECT_EQ(lac.cachedResourceGroupCount(), 0); + ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); + ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id + 1, "analytics")); + + ASSERT_EQ(request_has_keyspace.size(), 2); + EXPECT_TRUE(request_has_keyspace[0]); + EXPECT_TRUE(request_has_keyspace[1]); + + EXPECT_EQ(lac.getCachedResourceGroupForTest(NullspaceID, "default"), nullptr); + auto default_group = lac.getCachedResourceGroupForTest(keyspace_id, "default"); + ASSERT_NE(default_group, nullptr); + auto request_info = default_group->buildRequestInfoIfNecessary(SteadyClock::now()); + ASSERT_TRUE(request_info.has_value()); + EXPECT_EQ(request_info->keyspace_id, keyspace_id); + + auto analytics_group = lac.getCachedResourceGroupForTest(keyspace_id + 1, "analytics"); + ASSERT_NE(analytics_group, nullptr); + EXPECT_EQ(analytics_group->group_pb.keyspace_id().value(), keyspace_id + 1); } -TEST(LocalAdmissionControllerTest, ReservedDefaultGroupStillFallbacksWhenLegacyLookupThrows) +TEST(LocalAdmissionControllerTest, WarmupErrorPropagatesWithoutCompatibilityFallback) { constexpr KeyspaceID keyspace_id = 31; + size_t request_count = 0; pingcap::kv::Cluster cluster; - cluster.pd_client = std::make_shared([](const resource_manager::GetResourceGroupRequest & req) { - if (!req.has_keyspace_id()) - throw pingcap::Exception("legacy lookup failed", pingcap::ErrorCodes::UnknownError); - + cluster.pd_client = std::make_shared([&request_count](const resource_manager::GetResourceGroupRequest &) { + ++request_count; resource_manager::GetResourceGroupResponse resp; resp.mutable_error()->set_message("resource group default not found"); return resp; }); - LocalAdmissionController lac(&cluster, nullptr, false, false); + LocalAdmissionController lac(&cluster, nullptr, true, false); - ASSERT_NO_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default")); - auto group = lac.getCachedResourceGroupForTest(NullspaceID, "default"); - ASSERT_NE(group, nullptr); - EXPECT_FALSE(group->enable_gac); - EXPECT_EQ(group->user_ru_per_sec, LocalAdmissionController::RESERVED_DEFAULT_RESOURCE_GROUP_RU_PER_SEC); + EXPECT_THROW(lac.warmupResourceGroupInfoCache(keyspace_id, "default"), DB::Exception); + EXPECT_EQ(request_count, 1); + EXPECT_EQ(lac.cachedResourceGroupCount(), 0); } } // namespace DB::tests From 252479897b1aaae06a2d2a31006b818069df9772 Mon Sep 17 00:00:00 2001 From: yongman Date: Fri, 3 Jul 2026 21:51:51 +0800 Subject: [PATCH 7/7] format Signed-off-by: yongman --- .../ResourceControl/LocalAdmissionController.cpp | 13 ++++++++----- .../ResourceControl/LocalAdmissionController.h | 12 +++++------- .../tests/gtest_local_admission_controller.cpp | 13 +++++++------ 3 files changed, 20 insertions(+), 18 deletions(-) diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp index cc26aaee534..0205a9eb617 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.cpp @@ -351,14 +351,14 @@ void LocalAdmissionController::warmupResourceGroupInfoCache(const KeyspaceID & k resp.error().message()); checkGACRespValid(resp.group()); - const auto detected_mode - = with_keyspace ? detectBackendMode(resp.group()) : ResourceGroupBackendMode::LegacyGlobal; + const auto detected_mode = with_keyspace ? detectBackendMode(resp.group()) : ResourceGroupBackendMode::LegacyGlobal; // pd-cse returns a group without keyspace_id and keeps using the shared legacy namespace. // New PD returns the concrete keyspace-scoped definition. TiFlash only needs this one bit to // choose its cache key shape and which watch prefix to keep following afterwards. updateBackendMode(detected_mode); - const auto resolved_keyspace_id - = detected_mode == ResourceGroupBackendMode::LegacyGlobal ? NullspaceID : resolveResourceGroupKeyspace(resp.group()); + const auto resolved_keyspace_id = detected_mode == ResourceGroupBackendMode::LegacyGlobal + ? NullspaceID + : resolveResourceGroupKeyspace(resp.group()); addResourceGroup(resolved_keyspace_id, resp.group()); } @@ -935,7 +935,10 @@ bool LocalAdmissionController::handlePutEvent( if (rg == nullptr) { // Ignore updates for groups that have never been warmed into the local cache yet. - LOG_DEBUG(log, "trying to modify resource group config({}), but cannot find its info", group_pb.ShortDebugString()); + LOG_DEBUG( + log, + "trying to modify resource group config({}), but cannot find its info", + group_pb.ShortDebugString()); return true; } const auto old_user_ru_per_sec = rg->user_ru_per_sec; diff --git a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h index faa5d1b71b0..8f653c85a98 100644 --- a/dbms/src/Flash/ResourceControl/LocalAdmissionController.h +++ b/dbms/src/Flash/ResourceControl/LocalAdmissionController.h @@ -371,7 +371,8 @@ class LocalAdmissionController final : private boost::noncopyable , etcd_client(etcd_client_) , with_keyspace(with_keyspace) , start_background_threads(start_background_threads) - , resource_group_backend_mode(with_keyspace ? ResourceGroupBackendMode::Unknown : ResourceGroupBackendMode::LegacyGlobal) + , resource_group_backend_mode( + with_keyspace ? ResourceGroupBackendMode::Unknown : ResourceGroupBackendMode::LegacyGlobal) { current_tick = SteadyClock::now(); last_clear_cpu_time = current_tick; @@ -577,7 +578,8 @@ class LocalAdmissionController final : private boost::noncopyable return nullptr; return iter->second; } - ResourceGroupPtr findResourceGroupWithCompatWithoutLock(const KeyspaceID & keyspace_id, const std::string & name) const + ResourceGroupPtr findResourceGroupWithCompatWithoutLock(const KeyspaceID & keyspace_id, const std::string & name) + const { if (auto group = findResourceGroupWithoutLock(keyspace_id, name); likely(group != nullptr)) return group; @@ -604,11 +606,7 @@ class LocalAdmissionController final : private boost::noncopyable if (iter != keyspace_resource_groups.end()) return; - LOG_INFO( - log, - "add new resource group, keyspace={} info: {}", - keyspace_id, - new_group_pb.ShortDebugString()); + LOG_INFO(log, "add new resource group, keyspace={} info: {}", keyspace_id, new_group_pb.ShortDebugString()); auto new_group = std::make_shared(keyspace_id, new_group_pb, current_tick); keyspace_resource_groups.insert({{keyspace_id, new_group_pb.name()}, new_group}); diff --git a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp index ef13ea89a8d..db35f25af8f 100644 --- a/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp +++ b/dbms/src/Flash/ResourceControl/tests/gtest_local_admission_controller.cpp @@ -139,12 +139,13 @@ TEST(LocalAdmissionControllerTest, WarmupErrorPropagatesWithoutCompatibilityFall constexpr KeyspaceID keyspace_id = 31; size_t request_count = 0; pingcap::kv::Cluster cluster; - cluster.pd_client = std::make_shared([&request_count](const resource_manager::GetResourceGroupRequest &) { - ++request_count; - resource_manager::GetResourceGroupResponse resp; - resp.mutable_error()->set_message("resource group default not found"); - return resp; - }); + cluster.pd_client + = std::make_shared([&request_count](const resource_manager::GetResourceGroupRequest &) { + ++request_count; + resource_manager::GetResourceGroupResponse resp; + resp.mutable_error()->set_message("resource group default not found"); + return resp; + }); LocalAdmissionController lac(&cluster, nullptr, true, false);