From 50fce404093626c4684f4de36874d8b766e24cb6 Mon Sep 17 00:00:00 2001 From: Eugene Ostroukhov Date: Fri, 31 Mar 2023 14:03:44 -0700 Subject: [PATCH] [PSM Interop] Add custom_lb test client (#32757) --- CMakeLists.txt | 1 + build_autogenerated.yaml | 2 + test/cpp/interop/BUILD | 19 ++ test/cpp/interop/rpc_behavior_lb_policy.cc | 193 ++++++++++++++++++ test/cpp/interop/rpc_behavior_lb_policy.h | 33 +++ test/cpp/interop/xds_interop_client.cc | 3 + .../tests/custom_lb_test.py | 2 + 7 files changed, 253 insertions(+) create mode 100644 test/cpp/interop/rpc_behavior_lb_policy.cc create mode 100644 test/cpp/interop/rpc_behavior_lb_policy.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 1890df768e4..16bd13e3038 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -24423,6 +24423,7 @@ add_executable(xds_interop_client ${_gRPC_PROTO_GENS_DIR}/src/proto/grpc/testing/xds/v3/percent.grpc.pb.h src/cpp/server/admin/admin_services.cc src/cpp/server/csds/csds.cc + test/cpp/interop/rpc_behavior_lb_policy.cc test/cpp/interop/xds_interop_client.cc third_party/googletest/googletest/src/gtest-all.cc third_party/googletest/googlemock/src/gmock-all.cc diff --git a/build_autogenerated.yaml b/build_autogenerated.yaml index f090139bd4a..5fb4563a859 100644 --- a/build_autogenerated.yaml +++ b/build_autogenerated.yaml @@ -13381,6 +13381,7 @@ targets: language: c++ headers: - src/cpp/server/csds/csds.h + - test/cpp/interop/rpc_behavior_lb_policy.h src: - src/proto/grpc/testing/empty.proto - src/proto/grpc/testing/messages.proto @@ -13391,6 +13392,7 @@ targets: - src/proto/grpc/testing/xds/v3/percent.proto - src/cpp/server/admin/admin_services.cc - src/cpp/server/csds/csds.cc + - test/cpp/interop/rpc_behavior_lb_policy.cc - test/cpp/interop/xds_interop_client.cc deps: - grpc++_reflection diff --git a/test/cpp/interop/BUILD b/test/cpp/interop/BUILD index 50008ec5f0b..48557290c58 100644 --- a/test/cpp/interop/BUILD +++ b/test/cpp/interop/BUILD @@ -244,6 +244,7 @@ grpc_cc_binary( "absl/flags:flag", ], deps = [ + ":rpc_behavior_lb_policy", "//:grpc++", "//:grpc++_reflection", "//:grpcpp_admin", @@ -446,3 +447,21 @@ grpc_cc_library( "//src/proto/grpc/testing:messages_proto", ], ) + +grpc_cc_library( + name = "rpc_behavior_lb_policy", + srcs = [ + "rpc_behavior_lb_policy.cc", + ], + hdrs = [ + "rpc_behavior_lb_policy.h", + ], + external_deps = [ + ], + language = "C++", + tags = ["nobuilder"], + deps = [ + "//:grpc", + "//:grpc++", + ], +) diff --git a/test/cpp/interop/rpc_behavior_lb_policy.cc b/test/cpp/interop/rpc_behavior_lb_policy.cc new file mode 100644 index 00000000000..9c40801e27d --- /dev/null +++ b/test/cpp/interop/rpc_behavior_lb_policy.cc @@ -0,0 +1,193 @@ +// +// +// Copyright 2023 gRPC authors. +// +// 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 "test/cpp/interop/rpc_behavior_lb_policy.h" + +#include "absl/strings/str_format.h" + +#include "src/core/lib/iomgr/pollset_set.h" +#include "src/core/lib/json/json_args.h" +#include "src/core/lib/json/json_object_loader.h" + +namespace grpc { +namespace testing { + +namespace { + +using grpc_core::CoreConfiguration; +using grpc_core::Json; +using grpc_core::JsonArgs; +using grpc_core::JsonLoaderInterface; +using grpc_core::LoadBalancingPolicy; +using grpc_core::OrphanablePtr; +using grpc_core::RefCountedPtr; + +constexpr absl::string_view kRpcBehaviorLbPolicyName = + "test.RpcBehaviorLoadBalancer"; + +constexpr absl::string_view kRpcBehaviorMetadataKey = "rpc-behavior"; + +class RpcBehaviorLbPolicyConfig : public LoadBalancingPolicy::Config { + public: + static JsonLoaderInterface* JsonLoader(const JsonArgs&) { + static const auto kJsonLoader = + grpc_core::JsonObjectLoader() + .Field("rpcBehavior", &RpcBehaviorLbPolicyConfig::rpc_behavior_) + .Finish(); + return kJsonLoader; + } + + absl::string_view rpc_behavior() const { return rpc_behavior_; } + + private: + absl::string_view name() const override { return kRpcBehaviorLbPolicyName; } + + std::string rpc_behavior_; +}; + +class RpcBehaviorLbPolicy : public LoadBalancingPolicy { + public: + explicit RpcBehaviorLbPolicy(Args args) + : LoadBalancingPolicy(std::move(args), /*initial_refcount=*/2) { + Args delegate_args; + delegate_args.work_serializer = work_serializer(); + delegate_args.args = channel_args(); + delegate_args.channel_control_helper = + std::make_unique(RefCountedPtr(this)); + delegate_ = + CoreConfiguration::Get().lb_policy_registry().CreateLoadBalancingPolicy( + "pick_first", std::move(delegate_args)); + grpc_pollset_set_add_pollset_set(delegate_->interested_parties(), + interested_parties()); + } + + ~RpcBehaviorLbPolicy() override = default; + + absl::string_view name() const override { return kRpcBehaviorLbPolicyName; } + + absl::Status UpdateLocked(UpdateArgs args) override { + RefCountedPtr config = std::move(args.config); + rpc_behavior_ = std::string(config->rpc_behavior()); + return delegate_->UpdateLocked(std::move(args)); + } + + void ExitIdleLocked() override { delegate_->ExitIdleLocked(); } + + void ResetBackoffLocked() override { delegate_->ResetBackoffLocked(); } + + private: + class Picker : public SubchannelPicker { + public: + Picker(RefCountedPtr delegate_picker, + absl::string_view rpc_behavior) + : delegate_picker_(std::move(delegate_picker)), + rpc_behavior_(rpc_behavior) {} + + PickResult Pick(PickArgs args) override { + char* rpc_behavior_copy = static_cast( + args.call_state->Alloc(rpc_behavior_.length() + 1)); + strcpy(rpc_behavior_copy, rpc_behavior_.c_str()); + args.initial_metadata->Add(kRpcBehaviorMetadataKey, rpc_behavior_copy); + // Do pick. + return delegate_picker_->Pick(args); + } + + private: + RefCountedPtr delegate_picker_; + std::string rpc_behavior_; + }; + + class Helper : public ChannelControlHelper { + public: + explicit Helper(RefCountedPtr parent) + : parent_(std::move(parent)) {} + + RefCountedPtr CreateSubchannel( + grpc_core::ServerAddress address, + const grpc_core::ChannelArgs& args) override { + return parent_->channel_control_helper()->CreateSubchannel( + std::move(address), args); + } + + void UpdateState(grpc_connectivity_state state, const absl::Status& status, + RefCountedPtr picker) override { + parent_->channel_control_helper()->UpdateState( + state, status, + grpc_core::MakeRefCounted(std::move(picker), + parent_->rpc_behavior_)); + } + + void RequestReresolution() override { + parent_->channel_control_helper()->RequestReresolution(); + } + + absl::string_view GetAuthority() override { + return parent_->channel_control_helper()->GetAuthority(); + } + + grpc_event_engine::experimental::EventEngine* GetEventEngine() override { + return parent_->channel_control_helper()->GetEventEngine(); + } + + void AddTraceEvent(TraceSeverity severity, + absl::string_view message) override { + parent_->channel_control_helper()->AddTraceEvent(severity, message); + } + + private: + RefCountedPtr parent_; + }; + + void ShutdownLocked() override { + grpc_pollset_set_del_pollset_set(delegate_->interested_parties(), + interested_parties()); + delegate_.reset(); + } + + OrphanablePtr delegate_; + std::string rpc_behavior_; +}; + +class RpcBehaviorLbPolicyFactory + : public grpc_core::LoadBalancingPolicyFactory { + private: + absl::string_view name() const override { return kRpcBehaviorLbPolicyName; } + + grpc_core::OrphanablePtr CreateLoadBalancingPolicy( + LoadBalancingPolicy::Args args) const override { + return grpc_core::MakeOrphanable(std::move(args)); + } + + absl::StatusOr> + ParseLoadBalancingConfig(const Json& json) const override { + return grpc_core::LoadRefCountedFromJson( + json, JsonArgs(), "errors validating LB policy config"); + } +}; +} // namespace + +void RegisterRpcBehaviorLbPolicy( + grpc_core::CoreConfiguration::Builder* builder) { + builder->lb_policy_registry()->RegisterLoadBalancingPolicyFactory( + std::make_unique()); +} + +} // namespace testing +} // namespace grpc diff --git a/test/cpp/interop/rpc_behavior_lb_policy.h b/test/cpp/interop/rpc_behavior_lb_policy.h new file mode 100644 index 00000000000..f0d6504b2f9 --- /dev/null +++ b/test/cpp/interop/rpc_behavior_lb_policy.h @@ -0,0 +1,33 @@ +// +// +// Copyright 2023 gRPC authors. +// +// 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 GRPC_TEST_CPP_INTEROP_RPC_BEHAVIOR_LB_POLICY_H +#define GRPC_TEST_CPP_INTEROP_RPC_BEHAVIOR_LB_POLICY_H + +#include + +#include "src/core/lib/config/core_configuration.h" + +namespace grpc { +namespace testing { +void RegisterRpcBehaviorLbPolicy( + grpc_core::CoreConfiguration::Builder* builder); +} // namespace testing +} // namespace grpc + +#endif // GRPC_TEST_CPP_INTEROP_RPC_BEHAVIOR_LB_POLICY_H \ No newline at end of file diff --git a/test/cpp/interop/xds_interop_client.cc b/test/cpp/interop/xds_interop_client.cc index f9105cb3edf..ecb3519860e 100644 --- a/test/cpp/interop/xds_interop_client.cc +++ b/test/cpp/interop/xds_interop_client.cc @@ -45,6 +45,7 @@ #include "src/proto/grpc/testing/messages.pb.h" #include "src/proto/grpc/testing/test.grpc.pb.h" #include "test/core/util/test_config.h" +#include "test/cpp/interop/rpc_behavior_lb_policy.h" #include "test/cpp/util/test_config.h" ABSL_FLAG(bool, fail_on_failed_rpc, false, @@ -580,6 +581,8 @@ void BuildRpcConfigsFromFlags(RpcConfigurationsQueue* rpc_configs_queue) { } int main(int argc, char** argv) { + grpc_core::CoreConfiguration::RegisterBuilder( + grpc::testing::RegisterRpcBehaviorLbPolicy); grpc::testing::TestEnvironment env(&argc, argv); grpc::testing::InitTest(&argc, &argv, true); // Validate the expect_status flag. diff --git a/tools/run_tests/xds_k8s_test_driver/tests/custom_lb_test.py b/tools/run_tests/xds_k8s_test_driver/tests/custom_lb_test.py index b09aed2f986..33c6ded105d 100644 --- a/tools/run_tests/xds_k8s_test_driver/tests/custom_lb_test.py +++ b/tools/run_tests/xds_k8s_test_driver/tests/custom_lb_test.py @@ -40,6 +40,8 @@ class CustomLbTest(xds_k8s_testcase.RegularXdsKubernetesTestCase): def is_supported(config: skips.TestConfig) -> bool: if config.client_lang == _Lang.JAVA: return config.version_gte('v1.47.x') + if config.client_lang == _Lang.CPP: + return config.version_gte('v1.55.x') return False def test_custom_lb_config(self):