From 7142a91fc961fa6be7093134b13d7b00896e40f0 Mon Sep 17 00:00:00 2001
From: vjpai <vpai@google.com>
Date: Thu, 23 Jun 2016 10:16:00 -0700
Subject: [PATCH] Fix up service types and method handlers so that FC unary can
 work properly.

---
 include/grpc++/impl/codegen/rpc_method.h         | 3 ++-
 include/grpc++/impl/codegen/rpc_service_method.h | 1 +
 include/grpc++/impl/codegen/service_type.h       | 7 +++++++
 src/compiler/cpp_generator.cc                    | 9 +++++++++
 4 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/include/grpc++/impl/codegen/rpc_method.h b/include/grpc++/impl/codegen/rpc_method.h
index 53bee3727ef..728ba3e334d 100644
--- a/include/grpc++/impl/codegen/rpc_method.h
+++ b/include/grpc++/impl/codegen/rpc_method.h
@@ -61,11 +61,12 @@ class RpcMethod {
 
   const char* name() const { return name_; }
   RpcType method_type() const { return method_type_; }
+  void SetMethodType(RpcType type) { method_type_ = type; }
   void* channel_tag() const { return channel_tag_; }
 
  private:
   const char* const name_;
-  const RpcType method_type_;
+  RpcType method_type_;
   void* const channel_tag_;
 };
 
diff --git a/include/grpc++/impl/codegen/rpc_service_method.h b/include/grpc++/impl/codegen/rpc_service_method.h
index 8b1f026c912..02da2a26177 100644
--- a/include/grpc++/impl/codegen/rpc_service_method.h
+++ b/include/grpc++/impl/codegen/rpc_service_method.h
@@ -82,6 +82,7 @@ class RpcServiceMethod : public RpcMethod {
   // if MethodHandler is nullptr, then this is an async method
   MethodHandler* handler() const { return handler_.get(); }
   void ResetHandler() { handler_.reset(); }
+  void SetHandler(MethodHandler *handler) { handler_.reset(handler); }
 
  private:
   void* server_tag_;
diff --git a/include/grpc++/impl/codegen/service_type.h b/include/grpc++/impl/codegen/service_type.h
index c19dfc7d45f..c0eeace8a4b 100644
--- a/include/grpc++/impl/codegen/service_type.h
+++ b/include/grpc++/impl/codegen/service_type.h
@@ -147,6 +147,13 @@ class Service {
     methods_[index].reset();
   }
 
+  void MarkMethodFCUnary(int index, MethodHandler* fc_unary_method) {
+    GPR_CODEGEN_ASSERT(methods_[index] && methods_[index]->handler() &&
+		       "Cannot mark an async or generic method as FCUnary");
+    methods_[index]->SetMethodType(::grpc::RpcMethod::FC_UNARY);
+    methods_[index]->SetHandler(fc_unary_method);
+  }
+
  private:
   friend class Server;
   friend class ServerInterface;
diff --git a/src/compiler/cpp_generator.cc b/src/compiler/cpp_generator.cc
index eadbb3be85c..bbd476783cc 100644
--- a/src/compiler/cpp_generator.cc
+++ b/src/compiler/cpp_generator.cc
@@ -129,6 +129,7 @@ grpc::string GetHeaderIncludes(File *file,
       "grpc++/impl/codegen/async_stream.h",
       "grpc++/impl/codegen/async_unary_call.h",
       "grpc++/impl/codegen/fc_unary.h",
+      "grpc++/impl/codegen/method_handler_impl.h",
       "grpc++/impl/codegen/proto_utils.h",
       "grpc++/impl/codegen/rpc_method.h",
       "grpc++/impl/codegen/service_type.h",
@@ -625,6 +626,11 @@ void PrintHeaderServerMethodFCUnary(
     printer->Indent();
     printer->Print(*vars,
 		   "WithFCUnaryMethod_$Method$() {\n"
+		   "  ::grpc::Status (*fn)(::grpc::ServerContext*, ::grpc::FCUnary< $Request$,$Response$>*) = this->WithFCUnaryMethod_$Method$<BaseClass>::$Method$;\n"
+		   "  ::grpc::Service::MarkMethodFCUnary($Idx$,\n"
+		   "    new ::grpc::FCUnaryMethodHandler<Service, "
+		   "$Request$, "
+		   "$Response$>(fn, this));\n"
 		   "}\n");
     printer->Print(*vars,
 		   "~WithFCUnaryMethod_$Method$() GRPC_OVERRIDE {\n"
@@ -1138,6 +1144,9 @@ void PrintSourceService(Printer *printer,
     (*vars)["Idx"] = as_string(i);
     if (method->NoStreaming()) {
       (*vars)["StreamingType"] = "NORMAL_RPC";
+      // NOTE: There is no reason to consider FC_UNARY as a separate
+      // category here since this part is setting up the client-side stub
+      // and this appears as a NORMAL_RPC from the client-side.
     } else if (method->ClientOnlyStreaming()) {
       (*vars)["StreamingType"] = "CLIENT_STREAMING";
     } else if (method->ServerOnlyStreaming()) {