Merge pull request #24145 from ananda1066/grpc_cli

Added timeout flag to grpc cli.

src/proto/grpc/testing/echo.proto

@@ -15,15 +15,17 @@
 
 syntax = "proto3";
 
+package grpc.testing;
+
 import "src/proto/grpc/testing/echo_messages.proto";
 import "src/proto/grpc/testing/simple_messages.proto";
 
-package grpc.testing;
-
 service EchoTestService {
   rpc Echo(EchoRequest) returns (EchoResponse);
   rpc Echo1(EchoRequest) returns (EchoResponse);
   rpc Echo2(EchoRequest) returns (EchoResponse);
+  rpc CheckDeadlineUpperBound(SimpleRequest) returns (StringValue);
+  rpc CheckDeadlineSet(SimpleRequest) returns (StringValue);
   // A service which checks that the initial metadata sent over contains some
   // expected key value pair
   rpc CheckClientInitialMetadata(SimpleRequest) returns (SimpleResponse);
@@ -64,5 +66,4 @@ service UnimplementedEchoService {
 }
 
 // A service without any rpc defined to test coverage.
-service NoRpcService {
-}
+service NoRpcService {}

src/proto/grpc/testing/simple_messages.proto

@@ -17,8 +17,10 @@ syntax = "proto3";
 
 package grpc.testing;
 
-message SimpleRequest {
-}
+message SimpleRequest {}
+
+message SimpleResponse {}
 
-message SimpleResponse {
+message StringValue {
+  string message = 1;
 }

test/cpp/util/cli_call.cc

@@ -18,9 +18,6 @@
 
 #include "test/cpp/util/cli_call.h"
 
-#include <iostream>
-#include <utility>
-
 #include <grpc/grpc.h>
 #include <grpc/slice.h>
 #include <grpc/support/log.h>
@@ -28,6 +25,10 @@
 #include <grpcpp/client_context.h>
 #include <grpcpp/support/byte_buffer.h>
 
+#include <cmath>
+#include <iostream>
+#include <utility>
+
 namespace grpc {
 namespace testing {
 namespace {
@@ -51,7 +52,7 @@ Status CliCall::Call(const std::shared_ptr<grpc::Channel>& channel,
 
 CliCall::CliCall(const std::shared_ptr<grpc::Channel>& channel,
                  const std::string& method,
-                 const OutgoingMetadataContainer& metadata)
+                 const OutgoingMetadataContainer& metadata, CliArgs args)
     : stub_(new grpc::GenericStub(channel)) {
   gpr_mu_init(&write_mu_);
   gpr_cv_init(&write_cv_);
@@ -61,6 +62,22 @@ CliCall::CliCall(const std::shared_ptr<grpc::Channel>& channel,
       ctx_.AddMetadata(iter->first, iter->second);
     }
   }
+
+  // Set deadline if timeout > 0 (default value -1 if no timeout specified)
+  if (args.timeout > 0) {
+    int64_t timeout_in_ns = ceil(args.timeout * 1e9);
+
+    // Convert timeout (in nanoseconds) to a deadline
+    auto deadline =
+        gpr_time_add(gpr_now(GPR_CLOCK_MONOTONIC),
+                     gpr_time_from_nanos(timeout_in_ns, GPR_TIMESPAN));
+    ctx_.set_deadline(deadline);
+  } else if (args.timeout != -1) {
+    fprintf(
+        stderr,
+        "WARNING: Non-positive timeout value, skipping setting deadline.\n");
+  }
+
   call_ = stub_->PrepareCall(&ctx_, method, &cq_);
   call_->StartCall(tag(1));
   void* got_tag;

test/cpp/util/cli_call.h

@@ -19,18 +19,22 @@
 #ifndef GRPC_TEST_CPP_UTIL_CLI_CALL_H
 #define GRPC_TEST_CPP_UTIL_CLI_CALL_H
 
-#include <map>
-
 #include <grpcpp/channel.h>
 #include <grpcpp/completion_queue.h>
 #include <grpcpp/generic/generic_stub.h>
 #include <grpcpp/support/status.h>
 #include <grpcpp/support/string_ref.h>
 
+#include <map>
+
 namespace grpc {
 
 class ClientContext;
 
+struct CliArgs {
+  double timeout = -1;
+};
+
 namespace testing {
 
 // CliCall handles the sending and receiving of generic messages given the name
@@ -43,7 +47,12 @@ class CliCall final {
       IncomingMetadataContainer;
 
   CliCall(const std::shared_ptr<grpc::Channel>& channel,
-          const std::string& method, const OutgoingMetadataContainer& metadata);
+          const std::string& method, const OutgoingMetadataContainer& metadata,
+          CliArgs args);
+  CliCall(const std::shared_ptr<grpc::Channel>& channel,
+          const std::string& method, const OutgoingMetadataContainer& metadata)
+      : CliCall(channel, method, metadata, CliArgs{}) {}
+
   ~CliCall();
 
   // Perform an unary generic RPC.

test/cpp/util/grpc_tool.cc

@@ -18,14 +18,6 @@
 
 #include "test/cpp/util/grpc_tool.h"
 
-#include <cstdio>
-#include <fstream>
-#include <iostream>
-#include <memory>
-#include <sstream>
-#include <string>
-#include <thread>
-
 #include <gflags/gflags.h>
 #include <grpc/grpc.h>
 #include <grpc/support/port_platform.h>
@@ -35,6 +27,14 @@
 #include <grpcpp/security/credentials.h>
 #include <grpcpp/support/string_ref.h>
 
+#include <cstdio>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <sstream>
+#include <string>
+#include <thread>
+
 #include "test/cpp/util/cli_call.h"
 #include "test/cpp/util/proto_file_parser.h"
 #include "test/cpp/util/proto_reflection_descriptor_database.h"
@@ -74,6 +74,9 @@ DEFINE_bool(batch, false,
             "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.");
+DEFINE_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.");
 
 namespace {
 
@@ -490,7 +493,10 @@ bool GrpcTool::CallMethod(int argc, const char** argv,
       "    --binary_input           ; Input in binary format\n"
       "    --binary_output          ; Output in binary format\n"
       "    --json_input             ; Input in json format\n"
-      "    --json_output            ; Output in json format\n" +
+      "    --json_output            ; Output in json format\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" +
       cred.GetCredentialUsage());
 
   std::stringstream output_ss;
@@ -500,6 +506,8 @@ bool GrpcTool::CallMethod(int argc, const char** argv,
   std::string formatted_method_name;
   std::unique_ptr<ProtoFileParser> parser;
   std::string serialized_request_proto;
+  CliArgs cli_args;
+  cli_args.timeout = FLAGS_timeout;
   bool print_mode = false;
 
   std::shared_ptr<grpc::Channel> channel =
@@ -544,7 +552,7 @@ bool GrpcTool::CallMethod(int argc, const char** argv,
     ParseMetadataFlag(&client_metadata);
     PrintMetadata(client_metadata, "Sending client initial metadata:");
 
-    CliCall call(channel, formatted_method_name, client_metadata);
+    CliCall call(channel, formatted_method_name, client_metadata, cli_args);
     if (FLAGS_display_peer_address) {
       fprintf(stderr, "New call for method_name:%s has peer address:|%s|\n",
               formatted_method_name.c_str(), call.peer().c_str());
@@ -677,7 +685,8 @@ bool GrpcTool::CallMethod(int argc, const char** argv,
           std::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);
+          CliCall call(channel, formatted_method_name, client_metadata,
+                       cli_args);
           if (FLAGS_display_peer_address) {
             fprintf(stderr,
                     "New call for method_name:%s has peer address:|%s|\n",
@@ -780,7 +789,7 @@ bool GrpcTool::CallMethod(int argc, const char** argv,
     ParseMetadataFlag(&client_metadata);
     PrintMetadata(client_metadata, "Sending client initial metadata:");
 
-    CliCall call(channel, formatted_method_name, client_metadata);
+    CliCall call(channel, formatted_method_name, client_metadata, cli_args);
     if (FLAGS_display_peer_address) {
       fprintf(stderr, "New call for method_name:%s has peer address:|%s|\n",
               formatted_method_name.c_str(), call.peer().c_str());

test/cpp/util/grpc_tool_test.cc

@@ -30,6 +30,7 @@
 #include <grpcpp/server_context.h>
 #include <gtest/gtest.h>
 
+#include <chrono>
 #include <sstream>
 
 #include "src/core/lib/gpr/env.h"
@@ -54,6 +55,8 @@ using grpc::testing::EchoResponse;
   "Echo\n"                        \
   "Echo1\n"                       \
   "Echo2\n"                       \
+  "CheckDeadlineUpperBound\n"     \
+  "CheckDeadlineSet\n"            \
   "CheckClientInitialMetadata\n"  \
   "RequestStream\n"               \
   "ResponseStream\n"              \
@@ -70,6 +73,10 @@ using grpc::testing::EchoResponse;
   "{}\n"                                                                       \
   "  rpc Echo2(grpc.testing.EchoRequest) returns (grpc.testing.EchoResponse) " \
   "{}\n"                                                                       \
+  "  rpc CheckDeadlineUpperBound(grpc.testing.SimpleRequest) returns "         \
+  "(grpc.testing.StringValue) {}\n"                                            \
+  "  rpc CheckDeadlineSet(grpc.testing.SimpleRequest) returns "                \
+  "(grpc.testing.StringValue) {}\n"                                            \
   "  rpc CheckClientInitialMetadata(grpc.testing.SimpleRequest) returns "      \
   "(grpc.testing.SimpleResponse) {}\n"                                         \
   "  rpc RequestStream(stream grpc.testing.EchoRequest) returns "              \
@@ -119,6 +126,7 @@ DECLARE_string(metadata);
 DECLARE_string(protofiles);
 DECLARE_string(proto_path);
 DECLARE_string(default_service_config);
+DECLARE_double(timeout);
 
 namespace {
 
@@ -177,6 +185,29 @@ class TestServiceImpl : public ::grpc::testing::EchoTestService::Service {
     return Status::OK;
   }
 
+  Status CheckDeadlineSet(ServerContext* context, const SimpleRequest* request,
+                          StringValue* response) override {
+    response->set_message(context->deadline() !=
+                                  std::chrono::system_clock::time_point::max()
+                              ? "true"
+                              : "false");
+    return Status::OK;
+  }
+
+  // Check if deadline - current time <= timeout
+  // If deadline set, timeout + current time should be an upper bound for it
+  Status CheckDeadlineUpperBound(ServerContext* context,
+                                 const SimpleRequest* request,
+                                 StringValue* response) override {
+    auto seconds = std::chrono::duration_cast<std::chrono::seconds>(
+        context->deadline() - std::chrono::system_clock::now());
+
+    // Returning string instead of bool to avoid using embedded messages in
+    // proto3
+    response->set_message(seconds.count() <= FLAGS_timeout ? "true" : "false");
+    return Status::OK;
+  }
+
   Status RequestStream(ServerContext* context,
                        ServerReader<EchoRequest>* reader,
                        EchoResponse* response) override {
@@ -862,6 +893,93 @@ TEST_F(GrpcToolTest, CallCommandRequestStreamWithBadRequestJsonInput) {
   ShutdownServer();
 }
 
+TEST_F(GrpcToolTest, CallCommandWithTimeoutDeadlineSet) {
+  // Test input "grpc_cli call CheckDeadlineSet --timeout=5000.25"
+  std::stringstream output_stream;
+
+  const std::string server_address = SetUpServer();
+  const char* argv[] = {"grpc_cli", "call", server_address.c_str(),
+                        "CheckDeadlineSet"};
+
+  // Set timeout to 5000.25 seconds
+  FLAGS_timeout = 5000.25;
+
+  EXPECT_TRUE(0 == GrpcToolMainLib(ArraySize(argv), argv, TestCliCredentials(),
+                                   std::bind(PrintStream, &output_stream,
+                                             std::placeholders::_1)));
+
+  // Expected output: "message: "true"", deadline set
+  EXPECT_TRUE(nullptr !=
+              strstr(output_stream.str().c_str(), "message: \"true\""));
+  ShutdownServer();
+}
+
+TEST_F(GrpcToolTest, CallCommandWithTimeoutDeadlineUpperBound) {
+  // Test input "grpc_cli call CheckDeadlineUpperBound --timeout=900"
+  std::stringstream output_stream;
+
+  const std::string server_address = SetUpServer();
+  const char* argv[] = {"grpc_cli", "call", server_address.c_str(),
+                        "CheckDeadlineUpperBound"};
+
+  // Set timeout to 900 seconds
+  FLAGS_timeout = 900;
+
+  EXPECT_TRUE(0 == GrpcToolMainLib(ArraySize(argv), argv, TestCliCredentials(),
+                                   std::bind(PrintStream, &output_stream,
+                                             std::placeholders::_1)));
+
+  // Expected output: "message: "true""
+  // deadline not greater than timeout + current time
+  EXPECT_TRUE(nullptr !=
+              strstr(output_stream.str().c_str(), "message: \"true\""));
+  ShutdownServer();
+}
+
+TEST_F(GrpcToolTest, CallCommandWithNegativeTimeoutValue) {
+  // Test input "grpc_cli call CheckDeadlineSet --timeout=-5"
+  std::stringstream output_stream;
+
+  const std::string server_address = SetUpServer();
+  const char* argv[] = {"grpc_cli", "call", server_address.c_str(),
+                        "CheckDeadlineSet"};
+
+  // Set timeout to -5 (deadline not set)
+  FLAGS_timeout = -5;
+
+  EXPECT_TRUE(0 == GrpcToolMainLib(ArraySize(argv), argv, TestCliCredentials(),
+                                   std::bind(PrintStream, &output_stream,
+                                             std::placeholders::_1)));
+
+  // Expected output: "message: "false"", deadline not set
+  EXPECT_TRUE(nullptr !=
+              strstr(output_stream.str().c_str(), "message: \"false\""));
+
+  ShutdownServer();
+}
+
+TEST_F(GrpcToolTest, CallCommandWithDefaultTimeoutValue) {
+  // Test input "grpc_cli call CheckDeadlineSet --timeout=-1"
+  std::stringstream output_stream;
+
+  const std::string server_address = SetUpServer();
+  const char* argv[] = {"grpc_cli", "call", server_address.c_str(),
+                        "CheckDeadlineSet"};
+
+  // Set timeout to -1 (default value, deadline not set)
+  FLAGS_timeout = -1;
+
+  EXPECT_TRUE(0 == GrpcToolMainLib(ArraySize(argv), argv, TestCliCredentials(),
+                                   std::bind(PrintStream, &output_stream,
+                                             std::placeholders::_1)));
+
+  // Expected output: "message: "false"", deadline not set
+  EXPECT_TRUE(nullptr !=
+              strstr(output_stream.str().c_str(), "message: \"false\""));
+
+  ShutdownServer();
+}
+
 TEST_F(GrpcToolTest, CallCommandResponseStream) {
   // Test input: grpc_cli call localhost:<port> ResponseStream "message:
   // 'Hello'"