diff options
Diffstat (limited to 'src')
3 files changed, 185 insertions, 1 deletions
diff --git a/src/buildtool/execution_api/remote/TARGETS b/src/buildtool/execution_api/remote/TARGETS index f5923306..fe04bbe6 100644 --- a/src/buildtool/execution_api/remote/TARGETS +++ b/src/buildtool/execution_api/remote/TARGETS @@ -4,6 +4,7 @@ , "hdrs": [ "bazel/bytestream_client.hpp" , "bazel/bazel_action.hpp" + , "bazel/bazel_capabilities_client.hpp" , "bazel/bazel_response.hpp" , "bazel/bazel_network.hpp" , "bazel/bazel_ac_client.hpp" @@ -13,6 +14,7 @@ ] , "srcs": [ "bazel/bazel_action.cpp" + , "bazel/bazel_capabilities_client.cpp" , "bazel/bazel_response.cpp" , "bazel/bazel_network.cpp" , "bazel/bazel_ac_client.cpp" @@ -33,6 +35,7 @@ , ["src/buildtool/execution_api/common", "artifact_blob"] , ["src/buildtool/execution_api/common", "bytestream_utils"] , ["src/buildtool/execution_api/common", "common"] + , ["src/buildtool/execution_api/common", "message_limits"] , ["src/buildtool/file_system", "git_repo"] , ["src/buildtool/logging", "log_level"] , ["src/buildtool/logging", "logging"] @@ -55,7 +58,6 @@ , ["src/buildtool/common/remote", "retry"] , ["src/buildtool/execution_api/bazel_msg", "bazel_msg_factory"] , ["src/buildtool/execution_api/common", "common_api"] - , ["src/buildtool/execution_api/common", "message_limits"] , ["src/buildtool/execution_api/utils", "outputscheck"] , ["src/buildtool/file_system", "object_type"] , ["src/utils/cpp", "back_map"] diff --git a/src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.cpp b/src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.cpp new file mode 100644 index 00000000..7a7968bd --- /dev/null +++ b/src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.cpp @@ -0,0 +1,116 @@ +// Copyright 2025 Huawei Cloud Computing Technology Co., Ltd. +// +// 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 "src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.hpp" + +#include <algorithm> +#include <mutex> +#include <optional> +#include <utility> + +#include <grpcpp/grpcpp.h> + +#include "fmt/core.h" +#include "src/buildtool/common/bazel_types.hpp" +#include "src/buildtool/common/remote/client_common.hpp" +#include "src/buildtool/common/remote/retry.hpp" +#include "src/buildtool/logging/log_level.hpp" + +namespace { +[[nodiscard]] auto Parse(std::optional<bazel_re::ServerCapabilities> + response) noexcept -> Capabilities { + if (not response.has_value()) { + return Capabilities{}; + } + + // To not duplicate default values here, create default capabilities and + // copy data from there. + Capabilities const default_capabilities; + + // If capabilities don't contain cache capabilities, max_batch_total_size is + // unlimited(equals 0) or greater than the internal limit, fall back to the + // default max_batch_total_size: + std::size_t max_batch = default_capabilities.MaxBatchTransferSize; + if (response->has_cache_capabilities() and + response->cache_capabilities().max_batch_total_size_bytes() != 0) { + max_batch = std::min<std::size_t>( + response->cache_capabilities().max_batch_total_size_bytes(), + default_capabilities.MaxBatchTransferSize); + } + return Capabilities{ + .MaxBatchTransferSize = max_batch, + }; +} +} // namespace + +BazelCapabilitiesClient::BazelCapabilitiesClient( + std::string const& server, + Port port, + gsl::not_null<Auth const*> const& auth, + gsl::not_null<RetryConfig const*> const& retry_config) noexcept + : retry_config_{*retry_config} { + stub_ = bazel_re::Capabilities::NewStub( + CreateChannelWithCredentials(server, port, auth)); +} + +auto BazelCapabilitiesClient::GetCapabilities( + std::string const& instance_name) const noexcept -> Capabilities::Ptr { + { + // Check the cache already contains capabilities for this instance: + std::shared_lock guard{lock_}; + auto it = capabilities_.find(instance_name); + if (it != capabilities_.end()) { + return it->second; + } + } + + std::optional<bazel_re::ServerCapabilities> response; + auto get_capabilities = + [&instance_name, &stub = *stub_, &response]() -> RetryResponse { + grpc::ClientContext context; + + bazel_re::GetCapabilitiesRequest request; + *request.mutable_instance_name() = instance_name; + + bazel_re::ServerCapabilities capabilities; + auto const status = + stub.GetCapabilities(&context, request, &capabilities); + if (status.ok()) { + response.emplace(std::move(capabilities)); + return RetryResponse{.ok = true}; + } + return RetryResponse{ + .ok = false, + .exit_retry_loop = false, + .error_msg = fmt::format("While obtaining capabilities: {}", + status.error_message())}; + }; + + if (not WithRetry(get_capabilities, retry_config_, logger_) or + not response.has_value()) { + logger_.Emit( + LogLevel::Warning, + "Failed to obtain Capabilities. Falling back to default values."); + } + + bool const cache_result = response.has_value(); + auto result = std::make_shared<Capabilities>(Parse(std::move(response))); + + // Cache results only if they contain meaningful non-default capabilities: + if (cache_result) { + std::unique_lock lock{lock_}; + capabilities_.insert_or_assign(instance_name, result); + } + return result; +} diff --git a/src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.hpp b/src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.hpp new file mode 100644 index 00000000..b33bb2de --- /dev/null +++ b/src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.hpp @@ -0,0 +1,66 @@ +// Copyright 2025 Huawei Cloud Computing Technology Co., Ltd. +// +// 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. + +#ifndef INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_CAPABILITIES_CLIENT_HPP +#define INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_CAPABILITIES_CLIENT_HPP + +#include <cstddef> +#include <functional> +#include <memory> +#include <shared_mutex> +#include <string> +#include <unordered_map> + +#include "build/bazel/remote/execution/v2/remote_execution.grpc.pb.h" +#include "gsl/gsl" +#include "src/buildtool/auth/authentication.hpp" +#include "src/buildtool/common/remote/port.hpp" +#include "src/buildtool/common/remote/retry_config.hpp" +#include "src/buildtool/execution_api/common/message_limits.hpp" +#include "src/buildtool/logging/logger.hpp" + +namespace bazel_re = build::bazel::remote::execution::v2; + +struct Capabilities final { + using Ptr = gsl::not_null<std::shared_ptr<Capabilities>>; + + std::size_t const MaxBatchTransferSize = MessageLimits::kMaxGrpcLength; +}; + +class BazelCapabilitiesClient final { + public: + explicit BazelCapabilitiesClient( + std::string const& server, + Port port, + gsl::not_null<Auth const*> const& auth, + gsl::not_null<RetryConfig const*> const& retry_config) noexcept; + + /// \brief Obtain server capabilities for instance_name. + /// \return Capabilities corresponding to the given instance_name. Requested + /// capabilities are cached if a valid response is received from the server. + /// Otherwise, the default capabilities are returned and the caching step is + /// skipped to try again next time. + [[nodiscard]] auto GetCapabilities( + std::string const& instance_name) const noexcept -> Capabilities::Ptr; + + private: + RetryConfig const& retry_config_; + std::unique_ptr<bazel_re::Capabilities::Stub> stub_; + Logger logger_{"RemoteCapabilitiesClient"}; + + mutable std::shared_mutex lock_; + mutable std::unordered_map<std::string, Capabilities::Ptr> capabilities_; +}; + +#endif // INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_CAPABILITIES_CLIENT_HPP |