diff options
Diffstat (limited to 'test/cpp/util/grpc_tool.cc')
-rw-r--r-- | test/cpp/util/grpc_tool.cc | 125 |
1 files changed, 120 insertions, 5 deletions
diff --git a/test/cpp/util/grpc_tool.cc b/test/cpp/util/grpc_tool.cc index bb6f878020..cd6084ac6d 100644 --- a/test/cpp/util/grpc_tool.cc +++ b/test/cpp/util/grpc_tool.cc @@ -58,6 +58,11 @@ DEFINE_string(protofiles, "", "Name of the proto file."); DEFINE_bool(binary_input, false, "Input in binary format"); DEFINE_bool(binary_output, false, "Output in binary format"); DEFINE_string(infile, "", "Input file (default is stdin)"); +DEFINE_bool(batch, false, + "Input contains multiple requests. Please do not use this to send " + "more than a few RPCs. gRPC CLI has very different performance " + "characteristics compared with normal RPC calls which make it " + "unsuitable for loadtesting or significant production traffic."); namespace { @@ -460,12 +465,17 @@ bool GrpcTool::CallMethod(int argc, const char** argv, return false; } + if (argc == 3) { + request_text = argv[2]; + } + if (parser->IsStreaming(method_name, true /* is_request */)) { std::istream* input_stream; std::ifstream input_file; - if (argc == 3) { - request_text = argv[2]; + if (FLAGS_batch) { + fprintf(stderr, "Batch mode for streaming RPC is not supported.\n"); + return false; } std::multimap<grpc::string, grpc::string> client_metadata; @@ -549,8 +559,115 @@ bool GrpcTool::CallMethod(int argc, const char** argv, } } else { // parser->IsStreaming(method_name, true /* is_request */) + if (FLAGS_batch) { + if (parser->IsStreaming(method_name, false /* is_request */)) { + fprintf(stderr, "Batch mode for streaming RPC is not supported.\n"); + return false; + } + + std::istream* input_stream; + std::ifstream input_file; + + if (FLAGS_infile.empty()) { + if (isatty(fileno(stdin))) { + print_mode = true; + fprintf(stderr, "reading request messages from stdin...\n"); + } + input_stream = &std::cin; + } else { + input_file.open(FLAGS_infile, std::ios::in | std::ios::binary); + input_stream = &input_file; + } + + std::multimap<grpc::string, grpc::string> client_metadata; + ParseMetadataFlag(&client_metadata); + if (print_mode) { + PrintMetadata(client_metadata, "Sending client initial metadata:"); + } + + std::stringstream request_ss; + grpc::string line; + while (!request_text.empty() || + (!input_stream->eof() && getline(*input_stream, line))) { + if (!request_text.empty()) { + if (FLAGS_binary_input) { + serialized_request_proto = request_text; + request_text.clear(); + } else { + serialized_request_proto = parser->GetSerializedProtoFromMethod( + method_name, request_text, true /* is_request */); + request_text.clear(); + if (parser->HasError()) { + if (print_mode) { + fprintf(stderr, "Failed to parse request.\n"); + } + continue; + } + } + + grpc::string serialized_response_proto; + std::multimap<grpc::string_ref, grpc::string_ref> + server_initial_metadata, server_trailing_metadata; + CliCall call(channel, formatted_method_name, client_metadata); + call.Write(serialized_request_proto); + call.WritesDone(); + if (!call.Read(&serialized_response_proto, + &server_initial_metadata)) { + fprintf(stderr, "Failed to read response.\n"); + } + Status status = call.Finish(&server_trailing_metadata); + + if (status.ok()) { + if (print_mode) { + fprintf(stderr, "Rpc succeeded with OK status.\n"); + PrintMetadata(server_initial_metadata, + "Received initial metadata from server:"); + PrintMetadata(server_trailing_metadata, + "Received trailing metadata from server:"); + } + + if (FLAGS_binary_output) { + if (!callback(serialized_response_proto)) { + break; + } + } else { + grpc::string response_text = parser->GetTextFormatFromMethod( + method_name, serialized_response_proto, + false /* is_request */); + if (parser->HasError() && print_mode) { + fprintf(stderr, "Failed to parse response.\n"); + } else { + if (!callback(response_text)) { + break; + } + } + } + } else { + if (print_mode) { + fprintf(stderr, + "Rpc failed with status code %d, error message: %s\n", + status.error_code(), status.error_message().c_str()); + } + } + } else { + if (line.length() == 0) { + request_text = request_ss.str(); + request_ss.str(grpc::string()); + request_ss.clear(); + } else { + request_ss << line << ' '; + } + } + } + + if (input_file.is_open()) { + input_file.close(); + } + + return true; + } + if (argc == 3) { - request_text = argv[2]; if (!FLAGS_infile.empty()) { fprintf(stderr, "warning: request given in argv, ignoring --infile\n"); } @@ -571,9 +688,7 @@ bool GrpcTool::CallMethod(int argc, const char** argv, if (FLAGS_binary_input) { serialized_request_proto = request_text; - // formatted_method_name = method_name; } else { - // formatted_method_name = parser->GetFormattedMethodName(method_name); serialized_request_proto = parser->GetSerializedProtoFromMethod( method_name, request_text, true /* is_request */); if (parser->HasError()) { |