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

@@ -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<grpc::Channel> CreateCliChannel(
-    const std::string& server_address, const CliCredentials& cred) {
+    const std::string& server_address, const CliCredentials& cred,
-  grpc::ChannelArguments args;
+    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<grpc::Channel> 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<grpc::Channel> 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<grpc::Channel> 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<grpc::Channel> channel =
-        CreateCliChannel(server_address, cred);
+        CreateCliChannel(server_address, cred, grpc::ChannelArguments());
     parser = absl::make_unique<grpc::testing::ProtoFileParser>(
         absl::GetFlag(FLAGS_remotedb) ? channel : nullptr,
         absl::GetFlag(FLAGS_proto_path), absl::GetFlag(FLAGS_protofiles));

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);