libbrillo: Add rpc_deadline arg to CallRpc of AsyncGrpcClientupstream-master2

BUG=b:151695024
TEST=TEST=cros_run_unit_tests --board $BOARD --packages \
     "libbrillo"

Change-Id: I868a9139da7be3c3bb3375727e79c6211743b564
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform2/+/2270357
Reviewed-by: Maksim Ivanov <emaxx@chromium.org>
Reviewed-by: Oleh Lamzin <lamzin@google.com>
Tested-by: Woody Chow <woodychow@chromium.org>
Commit-Queue: Woody Chow <woodychow@chromium.org>
Cr-Mirrored-From: https://chromium.googlesource.com/chromiumos/platform2
Cr-Mirrored-Commit: 3ac0784a15f53fa85f0d776f4ba75bf7588bbeda

4 files changed, 51 insertions, 12 deletions

diff --git a/brillo/grpc/async_grpc_client.h b/brillo/grpc/async_grpc_client.h
index 3359998..006607c 100644
--- a/brillo/grpc/async_grpc_client.h
+++ b/brillo/grpc/async_grpc_client.h
@@ -103,18 +103,19 @@ class AsyncGrpcClient final : public internal::AsyncGrpcClientBase {
                                const RequestType& request,
                                grpc::CompletionQueue* cq);
 
-  // Call RPC represented by |async_rpc_start|. Pass |request| as the
-  // request. Call |on_reply_callback| on the task runner passed to the
-  // constructor when a response is available.
+  // Call RPC represented by |async_rpc_start|. Pass |request| as the request
+  // with |rpc_deadline| as a timeout. Call |on_reply_callback| on the task
+  // runner passed to the constructor when a response is available.
   template <typename AsyncServiceStub,
             typename RequestType,
             typename ResponseType>
   void CallRpc(AsyncRequestFnPtr<AsyncServiceStub, RequestType, ResponseType>
                    async_rpc_start,
+               base::TimeDelta rpc_deadline,
                const RequestType& request,
                ReplyCallback<ResponseType> on_reply_callback) {
     std::unique_ptr<RpcState<ResponseType>> rpc_state =
-        std::make_unique<RpcState<ResponseType>>(rpc_deadline_);
+        std::make_unique<RpcState<ResponseType>>(rpc_deadline);
     RpcState<ResponseType>* rpc_state_unowned = rpc_state.get();
 
     std::unique_ptr<grpc::ClientAsyncResponseReader<ResponseType>> rpc =
@@ -132,9 +133,21 @@ class AsyncGrpcClient final : public internal::AsyncGrpcClientBase {
                 rpc_state_unowned->tag());
   }
 
+  // Same as above with the default deadline set by
+  // |SetDefaultRpcDeadlineForTesting|
+  template <typename AsyncServiceStub,
+            typename RequestType,
+            typename ResponseType>
+  void CallRpc(AsyncRequestFnPtr<AsyncServiceStub, RequestType, ResponseType>
+                   async_rpc_start,
+               const RequestType& request,
+               ReplyCallback<ResponseType> on_reply_callback) {
+    CallRpc(async_rpc_start, default_rpc_deadline_, request, on_reply_callback);
+  }
+
   // Sets the request deadline for future requests made with this client.
-  void SetRpcDeadlineForTesting(base::TimeDelta rpc_deadline) {
-    rpc_deadline_ = rpc_deadline;
+  void SetDefaultRpcDeadlineForTesting(base::TimeDelta default_rpc_deadline) {
+    default_rpc_deadline_ = default_rpc_deadline;
   }
 
  private:
@@ -172,7 +185,7 @@ class AsyncGrpcClient final : public internal::AsyncGrpcClientBase {
                           std::move(rpc_state->response));
   }
 
-  base::TimeDelta rpc_deadline_ = kRpcDeadline;
+  base::TimeDelta default_rpc_deadline_ = kDefaultRpcDeadline;
   std::unique_ptr<typename ServiceType::Stub> stub_;
 
   DISALLOW_COPY_AND_ASSIGN(AsyncGrpcClient);
diff --git a/brillo/grpc/async_grpc_client_server_test.cc b/brillo/grpc/async_grpc_client_server_test.cc
index b19fcb9..c75317a 100644
--- a/brillo/grpc/async_grpc_client_server_test.cc
+++ b/brillo/grpc/async_grpc_client_server_test.cc
@@ -600,7 +600,8 @@ TEST_F(AsyncGrpcClientServerTest, RpcServerRestarted) {
 TEST_F(AsyncGrpcClientServerTest, RpcServerStopped) {
   ShutDownServer();
 
-  client_->SetRpcDeadlineForTesting(base::TimeDelta::FromMilliseconds(50));
+  client_->SetDefaultRpcDeadlineForTesting(
+    base::TimeDelta::FromMilliseconds(50));
 
   base::TimeTicks start = base::TimeTicks::Now();
 
@@ -618,6 +619,31 @@ TEST_F(AsyncGrpcClientServerTest, RpcServerStopped) {
   EXPECT_GT(duration.InMilliseconds(), 40);  // Forgiving time comparison.
 }
 
+// Like RpcServerStopped. Pass context deadline to CallRpc directly.
+TEST_F(AsyncGrpcClientServerTest, RpcServerStopped_PerRequestTimeout) {
+  ShutDownServer();
+
+  client_->SetDefaultRpcDeadlineForTesting(
+    base::TimeDelta::FromMilliseconds(50));
+
+  base::TimeTicks start = base::TimeTicks::Now();
+
+  RpcReply<test_rpcs::EchoIntRpcResponse> rpc_reply;
+  test_rpcs::EchoIntRpcRequest request;
+  request.set_int_to_echo(1);
+  client_->CallRpc(&test_rpcs::ExampleService::Stub::AsyncEchoIntRpc,
+                   base::TimeDelta::FromMilliseconds(200),
+                   request,
+                   rpc_reply.MakeWriter());
+
+  rpc_reply.Wait();
+  EXPECT_TRUE(rpc_reply.IsError());
+
+  base::TimeDelta duration = base::TimeTicks::Now() - start;
+
+  EXPECT_GT(duration.InMilliseconds(), 180);  // Forgiving time comparison.
+}
+
 // Send a request to a server that starts after the request is made. The client
 // should only send the request after the connection has been established.
 TEST_F(AsyncGrpcClientServerTest, RpcServerStartedAfter) {
diff --git a/brillo/grpc/async_grpc_constants.cc b/brillo/grpc/async_grpc_constants.cc
index 4b7446c..1e92ee3 100644
--- a/brillo/grpc/async_grpc_constants.cc
+++ b/brillo/grpc/async_grpc_constants.cc
@@ -15,6 +15,6 @@ const base::TimeDelta kInitialGrpcReconnectBackoffTime =
 const base::TimeDelta kMaxGrpcReconnectBackoffTime =
     base::TimeDelta::FromSeconds(5);
 
-const base::TimeDelta kRpcDeadline = base::TimeDelta::FromMinutes(1);
+const base::TimeDelta kDefaultRpcDeadline = base::TimeDelta::FromMinutes(1);
 
 }  // namespace brillo
diff --git a/brillo/grpc/async_grpc_constants.h b/brillo/grpc/async_grpc_constants.h
index 55ba8be..969fbf9 100644
--- a/brillo/grpc/async_grpc_constants.h
+++ b/brillo/grpc/async_grpc_constants.h
@@ -26,9 +26,9 @@ BRILLO_EXPORT extern const base::TimeDelta kInitialGrpcReconnectBackoffTime;
 // Sets GRPC_ARG_MAX_RECONNECT_BACKOFF_MS
 BRILLO_EXPORT extern const base::TimeDelta kMaxGrpcReconnectBackoffTime;
 
-// Use this constant to set the deadline for RPC requests performed by the
-// GRPC client.
-BRILLO_EXPORT extern const base::TimeDelta kRpcDeadline;
+// Use this constant to set the default deadline for RPC requests performed by
+// the GRPC client.
+BRILLO_EXPORT extern const base::TimeDelta kDefaultRpcDeadline;
 
 }  // namespace brillo