From 747d016397ac52789e4c671c198a8b696be25021 Mon Sep 17 00:00:00 2001 From: Yousuk Seung Date: Fri, 23 Sep 2022 11:13:29 -0700 Subject: [PATCH] Support --max_recv_msg_size in grpc cli (#31106) * Support --max_recv_msg_size in grpc cli * Comment typos fixed in tests. --- test/cpp/util/grpc_tool.cc | 24 +++++++--- test/cpp/util/grpc_tool_test.cc | 85 +++++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 6 deletions(-) diff --git a/test/cpp/util/grpc_tool.cc b/test/cpp/util/grpc_tool.cc index 11986e6bd7e..d2238db2853 100644 --- a/test/cpp/util/grpc_tool.cc +++ b/test/cpp/util/grpc_tool.cc @@ -79,6 +79,10 @@ ABSL_FLAG(bool, batch, false, ABSL_FLAG(double, timeout, -1, "Specify timeout in seconds, used to set the deadline for all " "RPCs. The default value of -1 means no deadline has been set."); +ABSL_FLAG( + int, max_recv_msg_size, 0, + "Specify the max receive message size in bytes for all RPCs. -1 indicates " + "unlimited. The default value of 0 means to use the gRPC default."); namespace grpc { namespace testing { @@ -228,8 +232,9 @@ void ReadResponse(CliCall* call, const std::string& method_name, } std::shared_ptr CreateCliChannel( - const std::string& server_address, const CliCredentials& cred) { - grpc::ChannelArguments args; + const std::string& server_address, const CliCredentials& cred, + const grpc::ChannelArguments& extra_args) { + grpc::ChannelArguments args(extra_args); if (!cred.GetSslTargetNameOverride().empty()) { args.SetSslTargetNameOverride(cred.GetSslTargetNameOverride()); } @@ -363,7 +368,7 @@ bool GrpcTool::ListServices(int argc, const char** argv, std::string server_address(argv[0]); std::shared_ptr channel = - CreateCliChannel(server_address, cred); + CreateCliChannel(server_address, cred, grpc::ChannelArguments()); grpc::ProtoReflectionDescriptorDatabase desc_db(channel); grpc::protobuf::DescriptorPool desc_pool(&desc_db); @@ -464,7 +469,7 @@ bool GrpcTool::PrintType(int /*argc*/, const char** argv, std::string server_address(argv[0]); std::shared_ptr channel = - CreateCliChannel(server_address, cred); + CreateCliChannel(server_address, cred, grpc::ChannelArguments()); grpc::ProtoReflectionDescriptorDatabase desc_db(channel); grpc::protobuf::DescriptorPool desc_pool(&desc_db); @@ -504,6 +509,9 @@ bool GrpcTool::CallMethod(int argc, const char** argv, " --binary_output ; Output in binary format\n" " --json_input ; Input in json format\n" " --json_output ; Output in json format\n" + " --max_recv_msg_size ; Specify max receive message size in " + "bytes. -1 indicates unlimited. The default value of 0 means to use the " + "gRPC default.\n" " --timeout ; Specify timeout (in seconds), used to " "set the deadline for RPCs. The default value of -1 means no " "deadline has been set.\n" + @@ -520,8 +528,12 @@ bool GrpcTool::CallMethod(int argc, const char** argv, cli_args.timeout = absl::GetFlag(FLAGS_timeout); bool print_mode = false; + grpc::ChannelArguments args; + if (absl::GetFlag(FLAGS_max_recv_msg_size) != 0) { + args.SetMaxReceiveMessageSize(absl::GetFlag(FLAGS_max_recv_msg_size)); + } std::shared_ptr channel = - CreateCliChannel(server_address, cred); + CreateCliChannel(server_address, cred, args); if (!absl::GetFlag(FLAGS_binary_input) || !absl::GetFlag(FLAGS_binary_output)) { @@ -910,7 +922,7 @@ bool GrpcTool::ParseMessage(int argc, const char** argv, if (!absl::GetFlag(FLAGS_binary_input) || !absl::GetFlag(FLAGS_binary_output)) { std::shared_ptr channel = - CreateCliChannel(server_address, cred); + CreateCliChannel(server_address, cred, grpc::ChannelArguments()); parser = absl::make_unique( absl::GetFlag(FLAGS_remotedb) ? channel : nullptr, absl::GetFlag(FLAGS_proto_path), absl::GetFlag(FLAGS_protofiles)); diff --git a/test/cpp/util/grpc_tool_test.cc b/test/cpp/util/grpc_tool_test.cc index 2a927a606e8..85d5323d666 100644 --- a/test/cpp/util/grpc_tool_test.cc +++ b/test/cpp/util/grpc_tool_test.cc @@ -130,6 +130,7 @@ ABSL_DECLARE_FLAG(std::string, protofiles); ABSL_DECLARE_FLAG(std::string, proto_path); ABSL_DECLARE_FLAG(std::string, default_service_config); ABSL_DECLARE_FLAG(double, timeout); +ABSL_DECLARE_FLAG(int, max_recv_msg_size); namespace grpc { namespace testing { @@ -1296,6 +1297,90 @@ TEST_F(GrpcToolTest, CallCommandWithBadMetadata) { absl::SetFlag(&FLAGS_protofiles, ""); } +TEST_F(GrpcToolTest, CallMaxRecvMessageSizeSmall) { + std::stringstream output_stream; + const std::string server_address = SetUpServer(); + // Test input "grpc_cli call localhost:10000 Echo "message: 'Hello' + // --max_recv_msg_size=4" + const char* argv[] = {"grpc_cli", "call", server_address.c_str(), + "grpc.testing.EchoTestService.Echo", + "message: 'Hello'"}; + + // Set max_recv_msg_size to 4 which is not enough. + absl::SetFlag(&FLAGS_max_recv_msg_size, 4); + + // This should fail. + EXPECT_FALSE(0 == GrpcToolMainLib(ArraySize(argv), argv, TestCliCredentials(), + std::bind(PrintStream, &output_stream, + std::placeholders::_1))); + // No output expected. + EXPECT_TRUE(0 == output_stream.tellp()); + ShutdownServer(); +} + +TEST_F(GrpcToolTest, CallMaxRecvMessageSizeEnough) { + std::stringstream output_stream; + const std::string server_address = SetUpServer(); + // Test input "grpc_cli call localhost:10000 Echo "message: 'Hello' + // --max_recv_msg_size=1048576" + const char* argv[] = {"grpc_cli", "call", server_address.c_str(), + "grpc.testing.EchoTestService.Echo", + "message: 'Hello'"}; + + // Set max_recv_msg_size to a large enough number. + absl::SetFlag(&FLAGS_max_recv_msg_size, 1024 * 1024); + + EXPECT_TRUE(0 == GrpcToolMainLib(ArraySize(argv), argv, TestCliCredentials(), + std::bind(PrintStream, &output_stream, + std::placeholders::_1))); + // Expected output: "message: \"Hello\"" + EXPECT_TRUE(nullptr != + strstr(output_stream.str().c_str(), "message: \"Hello\"")); + ShutdownServer(); +} + +TEST_F(GrpcToolTest, CallMaxRecvMessageSizeDefault) { + std::stringstream output_stream; + const std::string server_address = SetUpServer(); + // Test input "grpc_cli call localhost:10000 Echo "message: 'Hello' + // --max_recv_msg_size=0" + const char* argv[] = {"grpc_cli", "call", server_address.c_str(), + "grpc.testing.EchoTestService.Echo", + "message: 'Hello'"}; + + // Set max_recv_msg_size to gRPC default, which should suffice. + absl::SetFlag(&FLAGS_max_recv_msg_size, 0); + + EXPECT_TRUE(0 == GrpcToolMainLib(ArraySize(argv), argv, TestCliCredentials(), + std::bind(PrintStream, &output_stream, + std::placeholders::_1))); + // Expected output: "message: \"Hello\"" + EXPECT_TRUE(nullptr != + strstr(output_stream.str().c_str(), "message: \"Hello\"")); + ShutdownServer(); +} + +TEST_F(GrpcToolTest, CallMaxRecvMessageSizeUnlimited) { + std::stringstream output_stream; + const std::string server_address = SetUpServer(); + // Test input "grpc_cli call localhost:10000 Echo "message: 'Hello' + // --max_recv_msg_size=-1" + const char* argv[] = {"grpc_cli", "call", server_address.c_str(), + "grpc.testing.EchoTestService.Echo", + "message: 'Hello'"}; + + // Set max_recv_msg_size to unlimited (-1), which should work. + absl::SetFlag(&FLAGS_max_recv_msg_size, -1); + + EXPECT_TRUE(0 == GrpcToolMainLib(ArraySize(argv), argv, TestCliCredentials(), + std::bind(PrintStream, &output_stream, + std::placeholders::_1))); + // Expected output: "message: \"Hello\"" + EXPECT_TRUE(nullptr != + strstr(output_stream.str().c_str(), "message: \"Hello\"")); + ShutdownServer(); +} + TEST_F(GrpcToolTest, ListCommand_OverrideSslHostName) { const std::string server_address = SetUpServer(true);