From 0911e489e3fe22e2ca5d7c927dac83358f2f05b7 Mon Sep 17 00:00:00 2001
From: Yash Tibrewal <yashkt@google.com>
Date: Thu, 15 Nov 2018 13:55:56 -0800
Subject: [PATCH] Add a method to check whether the message was received
 successfully

---
 include/grpcpp/impl/codegen/call_op_set.h            |  6 ++++--
 include/grpcpp/impl/codegen/interceptor.h            |  3 +++
 include/grpcpp/impl/codegen/interceptor_common.h     | 10 ++++++++++
 test/cpp/end2end/client_interceptors_end2end_test.cc | 12 ++++++++++++
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/include/grpcpp/impl/codegen/call_op_set.h b/include/grpcpp/impl/codegen/call_op_set.h
index aae8b9d3e37..3699ec94f2f 100644
--- a/include/grpcpp/impl/codegen/call_op_set.h
+++ b/include/grpcpp/impl/codegen/call_op_set.h
@@ -406,12 +406,13 @@ class CallOpRecvMessage {
 
   void SetInterceptionHookPoint(
       InterceptorBatchMethodsImpl* interceptor_methods) {
+    if (message_ == nullptr) return;
     interceptor_methods->SetRecvMessage(message_, &got_message);
   }
 
   void SetFinishInterceptionHookPoint(
       InterceptorBatchMethodsImpl* interceptor_methods) {
-    if (!got_message) return;
+    if (message_ == nullptr) return;
     interceptor_methods->AddInterceptionHookPoint(
         experimental::InterceptionHookPoints::POST_RECV_MESSAGE);
   }
@@ -501,12 +502,13 @@ class CallOpGenericRecvMessage {
 
   void SetInterceptionHookPoint(
       InterceptorBatchMethodsImpl* interceptor_methods) {
+    if (!deserialize_) return;
     interceptor_methods->SetRecvMessage(message_, &got_message);
   }
 
   void SetFinishInterceptionHookPoint(
       InterceptorBatchMethodsImpl* interceptor_methods) {
-    if (!got_message) return;
+    if (!deserialize_) return;
     interceptor_methods->AddInterceptionHookPoint(
         experimental::InterceptionHookPoints::POST_RECV_MESSAGE);
   }
diff --git a/include/grpcpp/impl/codegen/interceptor.h b/include/grpcpp/impl/codegen/interceptor.h
index 943376a5451..b977c350166 100644
--- a/include/grpcpp/impl/codegen/interceptor.h
+++ b/include/grpcpp/impl/codegen/interceptor.h
@@ -103,6 +103,9 @@ class InterceptorBatchMethods {
   // is already deserialized
   virtual void* GetRecvMessage() = 0;
 
+  // Checks whether the RECV MESSAGE op completed successfully
+  virtual bool GetRecvMessageStatus() = 0;
+
   // Returns a modifiable multimap of the received initial metadata
   virtual std::multimap<grpc::string_ref, grpc::string_ref>*
   GetRecvInitialMetadata() = 0;
diff --git a/include/grpcpp/impl/codegen/interceptor_common.h b/include/grpcpp/impl/codegen/interceptor_common.h
index d23b71f8a77..b2e92dd6f3d 100644
--- a/include/grpcpp/impl/codegen/interceptor_common.h
+++ b/include/grpcpp/impl/codegen/interceptor_common.h
@@ -103,6 +103,8 @@ class InterceptorBatchMethodsImpl
 
   void* GetRecvMessage() override { return recv_message_; }
 
+  bool GetRecvMessageStatus() override { return *got_message_; }
+
   std::multimap<grpc::string_ref, grpc::string_ref>* GetRecvInitialMetadata()
       override {
     return recv_initial_metadata_->map();
@@ -432,6 +434,14 @@ class CancelInterceptorBatchMethods
     return nullptr;
   }
 
+  bool GetRecvMessageStatus() override {
+    GPR_CODEGEN_ASSERT(
+        false &&
+        "It is illegal to call GetRecvMessageStatus on a method which "
+        "has a Cancel notification");
+    return false;
+  }
+
   std::multimap<grpc::string_ref, grpc::string_ref>* GetRecvInitialMetadata()
       override {
     GPR_CODEGEN_ASSERT(false &&
diff --git a/test/cpp/end2end/client_interceptors_end2end_test.cc b/test/cpp/end2end/client_interceptors_end2end_test.cc
index 459788ba519..4d4760a6524 100644
--- a/test/cpp/end2end/client_interceptors_end2end_test.cc
+++ b/test/cpp/end2end/client_interceptors_end2end_test.cc
@@ -325,6 +325,12 @@ class ServerStreamingRpcHijackingInterceptor
           static_cast<EchoResponse*>(methods->GetRecvMessage());
       resp->set_message("Hello");
     }
+    if (methods->QueryInterceptionHookPoint(
+            experimental::InterceptionHookPoints::POST_RECV_MESSAGE)) {
+      // Only the last message will be a failure
+      EXPECT_FALSE(got_failed_message_);
+      got_failed_message_ = !methods->GetRecvMessageStatus();
+    }
     if (methods->QueryInterceptionHookPoint(
             experimental::InterceptionHookPoints::PRE_RECV_STATUS)) {
       auto* map = methods->GetRecvTrailingMetadata();
@@ -341,11 +347,16 @@ class ServerStreamingRpcHijackingInterceptor
     }
   }
 
+  static bool GotFailedMessage() { return got_failed_message_; }
+
  private:
   experimental::ClientRpcInfo* info_;
+  static bool got_failed_message_;
   int count_ = 0;
 };
 
+bool ServerStreamingRpcHijackingInterceptor::got_failed_message_ = false;
+
 class ServerStreamingRpcHijackingInterceptorFactory
     : public experimental::ClientInterceptorFactoryInterface {
  public:
@@ -634,6 +645,7 @@ TEST_F(ClientInterceptorsStreamingEnd2endTest, ServerStreamingHijackingTest) {
   auto channel = experimental::CreateCustomChannelWithInterceptors(
       server_address_, InsecureChannelCredentials(), args, std::move(creators));
   MakeServerStreamingCall(channel);
+  EXPECT_TRUE(ServerStreamingRpcHijackingInterceptor::GotFailedMessage());
 }
 
 TEST_F(ClientInterceptorsStreamingEnd2endTest, BidiStreamingTest) {