4 files changed, 32 insertions, 6 deletions

diff --git a/include/grpcpp/impl/codegen/call_op_set.h b/include/grpcpp/impl/codegen/call_op_set.h
index 310bea93ca..b8a1c79ddf 100644
--- a/include/grpcpp/impl/codegen/call_op_set.h
+++ b/include/grpcpp/impl/codegen/call_op_set.h
@@ -436,14 +436,16 @@ class CallOpRecvMessage {
 
   void SetInterceptionHookPoint(
       InterceptorBatchMethodsImpl* interceptor_methods) {
-    interceptor_methods->SetRecvMessage(message_);
+    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);
+    if (!got_message) interceptor_methods->SetRecvMessage(nullptr, nullptr);
   }
   void SetHijackingState(InterceptorBatchMethodsImpl* interceptor_methods) {
     hijacked_ = true;
@@ -531,20 +533,23 @@ class CallOpGenericRecvMessage {
 
   void SetInterceptionHookPoint(
       InterceptorBatchMethodsImpl* interceptor_methods) {
-    interceptor_methods->SetRecvMessage(message_);
+    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);
+    if (!got_message) interceptor_methods->SetRecvMessage(nullptr, nullptr);
   }
   void SetHijackingState(InterceptorBatchMethodsImpl* interceptor_methods) {
     hijacked_ = true;
     if (!deserialize_) return;
     interceptor_methods->AddInterceptionHookPoint(
         experimental::InterceptionHookPoints::PRE_RECV_MESSAGE);
+    got_message = true;
   }
 
  private:
diff --git a/include/grpcpp/impl/codegen/interceptor.h b/include/grpcpp/impl/codegen/interceptor.h
index 5a9a3a44e6..9b49983748 100644
--- a/include/grpcpp/impl/codegen/interceptor.h
+++ b/include/grpcpp/impl/codegen/interceptor.h
@@ -162,6 +162,11 @@ class InterceptorBatchMethods {
   /// started from interceptors without infinite regress through the interceptor
   /// list.
   virtual std::unique_ptr<ChannelInterface> GetInterceptedChannel() = 0;
+
+  /// On a hijacked RPC, an interceptor can decide to fail a PRE_RECV_MESSAGE
+  /// op. This would be a signal to the reader that there will be no more
+  /// messages, or the stream has failed or been cancelled.
+  virtual void FailHijackedRecvMessage() = 0;
 };
 
 /// Interface for an interceptor. Interceptor authors must create a class
diff --git a/include/grpcpp/impl/codegen/interceptor_common.h b/include/grpcpp/impl/codegen/interceptor_common.h
index 4b7eaefee1..d60f9585fc 100644
--- a/include/grpcpp/impl/codegen/interceptor_common.h
+++ b/include/grpcpp/impl/codegen/interceptor_common.h
@@ -139,7 +139,10 @@ class InterceptorBatchMethodsImpl
     send_trailing_metadata_ = metadata;
   }
 
-  void SetRecvMessage(void* message) { recv_message_ = message; }
+  void SetRecvMessage(void* message, bool* got_message) {
+    recv_message_ = message;
+    got_message_ = got_message;
+  }
 
   void SetRecvInitialMetadata(MetadataMap* map) {
     recv_initial_metadata_ = map;
@@ -162,6 +165,12 @@ class InterceptorBatchMethodsImpl
         info->channel(), current_interceptor_index_ + 1));
   }
 
+  void FailHijackedRecvMessage() override {
+    GPR_CODEGEN_ASSERT(hooks_[static_cast<size_t>(
+        experimental::InterceptionHookPoints::PRE_RECV_MESSAGE)]);
+    *got_message_ = false;
+  }
+
   // Clears all state
   void ClearState() {
     reverse_ = false;
@@ -351,6 +360,7 @@ class InterceptorBatchMethodsImpl
   std::multimap<grpc::string, grpc::string>* send_trailing_metadata_ = nullptr;
 
   void* recv_message_ = nullptr;
+  bool* got_message_ = nullptr;
 
   MetadataMap* recv_initial_metadata_ = nullptr;
 
@@ -465,6 +475,12 @@ class CancelInterceptorBatchMethods
                        "method which has a Cancel notification");
     return std::unique_ptr<ChannelInterface>(nullptr);
   }
+
+  void FailHijackedRecvMessage() override {
+    GPR_CODEGEN_ASSERT(false &&
+                       "It is illegal to call FailHijackedRecvMessage on a "
+                       "method which has a Cancel notification");
+  }
 };
 }  // namespace internal
 }  // namespace grpc
diff --git a/include/grpcpp/impl/codegen/server_interface.h b/include/grpcpp/impl/codegen/server_interface.h
index e0e2629827..890a5650d0 100644
--- a/include/grpcpp/impl/codegen/server_interface.h
+++ b/include/grpcpp/impl/codegen/server_interface.h
@@ -272,7 +272,7 @@ class ServerInterface : public internal::CallHook {
       /* Set interception point for recv message */
       interceptor_methods_.AddInterceptionHookPoint(
           experimental::InterceptionHookPoints::POST_RECV_MESSAGE);
-      interceptor_methods_.SetRecvMessage(request_);
+      interceptor_methods_.SetRecvMessage(request_, nullptr);
       return RegisteredAsyncRequest::FinalizeResult(tag, status);
     }