diff --git a/include/pingcap/pd/Client.h b/include/pingcap/pd/Client.h index 16d9be57..2ebaa601 100644 --- a/include/pingcap/pd/Client.h +++ b/include/pingcap/pd/Client.h @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -64,6 +65,18 @@ class Client : public IClient std::string getLeaderUrl() override; + // ResourceControl related. + resource_manager::ListResourceGroupsResponse listResourceGroups(const resource_manager::ListResourceGroupsRequest &) override; + + resource_manager::GetResourceGroupResponse getResourceGroup(const resource_manager::GetResourceGroupRequest &) override; + + resource_manager::PutResourceGroupResponse addResourceGroup(const resource_manager::PutResourceGroupRequest &) override; + + resource_manager::PutResourceGroupResponse modifyResourceGroup(const resource_manager::PutResourceGroupRequest &) override; + + resource_manager::DeleteResourceGroupResponse deleteResourceGroup(const resource_manager::DeleteResourceGroupRequest &) override; + + std::shared_ptr> acquireTokenBuckets() override; private: void initClusterID(); @@ -82,6 +95,7 @@ class Client : public IClient std::shared_ptr channel; std::unique_ptr stub; std::unique_ptr keyspace_stub; + std::unique_ptr resource_manager_stub; PDConnClient(std::string addr, const ClusterConfig & config) { if (config.hasTlsConfig()) @@ -94,6 +108,7 @@ class Client : public IClient } stub = pdpb::PD::NewStub(channel); keyspace_stub = keyspacepb::Keyspace::NewStub(channel); + resource_manager_stub = resource_manager::ResourceManager::NewStub(channel); } }; diff --git a/include/pingcap/pd/IClient.h b/include/pingcap/pd/IClient.h index 8b634dab..7f62b0d9 100644 --- a/include/pingcap/pd/IClient.h +++ b/include/pingcap/pd/IClient.h @@ -9,6 +9,7 @@ #pragma GCC diagnostic ignored "-Wunused-parameter" #include #include +#include #pragma GCC diagnostic pop #include @@ -54,6 +55,19 @@ class IClient virtual bool isMock() = 0; virtual std::string getLeaderUrl() = 0; + + // ResourceControl related. + virtual resource_manager::ListResourceGroupsResponse listResourceGroups(const resource_manager::ListResourceGroupsRequest &) = 0; + + virtual resource_manager::GetResourceGroupResponse getResourceGroup(const resource_manager::GetResourceGroupRequest &) = 0; + + virtual resource_manager::PutResourceGroupResponse addResourceGroup(const resource_manager::PutResourceGroupRequest &) = 0; + + virtual resource_manager::PutResourceGroupResponse modifyResourceGroup(const resource_manager::PutResourceGroupRequest &) = 0; + + virtual resource_manager::DeleteResourceGroupResponse deleteResourceGroup(const resource_manager::DeleteResourceGroupRequest &) = 0; + + virtual std::shared_ptr> acquireTokenBuckets() = 0; }; using ClientPtr = std::shared_ptr; diff --git a/include/pingcap/pd/MockPDClient.h b/include/pingcap/pd/MockPDClient.h index 66140da8..04e8f43e 100644 --- a/include/pingcap/pd/MockPDClient.h +++ b/include/pingcap/pd/MockPDClient.h @@ -41,6 +41,24 @@ class MockPDClient : public IClient bool isMock() override { return true; } std::string getLeaderUrl() override { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } + + ::resource_manager::ListResourceGroupsResponse listResourceGroups(const ::resource_manager::ListResourceGroupsRequest &) override + { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } + + ::resource_manager::GetResourceGroupResponse getResourceGroup(const ::resource_manager::GetResourceGroupRequest &) override + { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } + + ::resource_manager::PutResourceGroupResponse addResourceGroup(const ::resource_manager::PutResourceGroupRequest &) override + { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } + + ::resource_manager::PutResourceGroupResponse modifyResourceGroup(const ::resource_manager::PutResourceGroupRequest &) override + { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } + + ::resource_manager::DeleteResourceGroupResponse deleteResourceGroup(const ::resource_manager::DeleteResourceGroupRequest &) override + { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } + + std::shared_ptr> acquireTokenBuckets() override + { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } }; } // namespace pd diff --git a/src/pd/Client.cc b/src/pd/Client.cc index fe710c0e..2b75fe23 100644 --- a/src/pd/Client.cc +++ b/src/pd/Client.cc @@ -508,5 +508,35 @@ bool Client::isClusterBootstrapped() return response.bootstrapped(); } +#define RESOURCE_CONTROL_FUNCTION_DEFINITION(FUNC_NAME, GRPC_METHOD, REQUEST_TYPE, RESPONSE_TYPE) \ + ::resource_manager::RESPONSE_TYPE Client::FUNC_NAME(const ::resource_manager::REQUEST_TYPE & request) \ + { \ + ::resource_manager::RESPONSE_TYPE response; \ + grpc::ClientContext context; \ + context.set_deadline(std::chrono::system_clock::now() + pd_timeout); \ + auto status = leaderClient()->resource_manager_stub->GRPC_METHOD(&context, request, &response); \ + if (!status.ok()) \ + { \ + std::string err_msg = ("resource manager grpc call failed: " #GRPC_METHOD ". " + std::to_string(status.error_code()) + ": " + status.error_message()); \ + log->error(err_msg); \ + check_leader.store(true); \ + throw Exception(err_msg, GRPCErrorCode); \ + } \ + return response; \ + } + +RESOURCE_CONTROL_FUNCTION_DEFINITION(listResourceGroups, ListResourceGroups, ListResourceGroupsRequest, ListResourceGroupsResponse) +RESOURCE_CONTROL_FUNCTION_DEFINITION(getResourceGroup, GetResourceGroup, GetResourceGroupRequest, GetResourceGroupResponse) +RESOURCE_CONTROL_FUNCTION_DEFINITION(addResourceGroup, AddResourceGroup, PutResourceGroupRequest, PutResourceGroupResponse) +RESOURCE_CONTROL_FUNCTION_DEFINITION(modifyResourceGroup, ModifyResourceGroup, PutResourceGroupRequest, PutResourceGroupResponse) +RESOURCE_CONTROL_FUNCTION_DEFINITION(deleteResourceGroup, DeleteResourceGroup, DeleteResourceGroupRequest, DeleteResourceGroupResponse) + +std::shared_ptr> Client::acquireTokenBuckets() +{ + resource_manager::TokenBucketsResponse response; + grpc::ClientContext context; + return leaderClient()->resource_manager_stub->AcquireTokenBuckets(&context); +} + } // namespace pd } // namespace pingcap