Populate channel tags for methods

5 files changed, 16 insertions, 19 deletions

diff --git a/include/grpc++/impl/internal_stub.h b/include/grpc++/impl/internal_stub.h
index 2cbf1d901b..1b223d034f 100644
--- a/include/grpc++/impl/internal_stub.h
+++ b/include/grpc++/impl/internal_stub.h
@@ -42,17 +42,13 @@ namespace grpc {
 
 class InternalStub {
  public:
-  InternalStub() {}
+  InternalStub(const std::shared_ptr<ChannelInterface>& channel) : channel_(channel) {}
   virtual ~InternalStub() {}
 
-  void set_channel(const std::shared_ptr<ChannelInterface>& channel) {
-    channel_ = channel;
-  }
-
   ChannelInterface* channel() { return channel_.get(); }
 
  private:
-  std::shared_ptr<ChannelInterface> channel_;
+  const std::shared_ptr<ChannelInterface> channel_;
 };
 
 }  // namespace grpc
diff --git a/include/grpc++/impl/rpc_method.h b/include/grpc++/impl/rpc_method.h
index 1346e0af06..3dfcdbad8f 100644
--- a/include/grpc++/impl/rpc_method.h
+++ b/include/grpc++/impl/rpc_method.h
@@ -45,15 +45,16 @@ class RpcMethod {
     BIDI_STREAMING
   };
 
-  RpcMethod(const char* name, RpcType type) : name_(name), method_type_(type) {}
+  RpcMethod(const char* name, RpcType type, void *channel_tag) : name_(name), method_type_(type), channel_tag_(channel_tag) {}
 
   const char* name() const { return name_; }
-
   RpcType method_type() const { return method_type_; }
+  void *channel_tag() const { return channel_tag_; }
 
  private:
-  const char* name_;
+  const char* const name_;
   const RpcType method_type_;
+  void * const channel_tag_;
 };
 
 }  // namespace grpc
diff --git a/include/grpc++/impl/rpc_service_method.h b/include/grpc++/impl/rpc_service_method.h
index 097667827a..50204d2099 100644
--- a/include/grpc++/impl/rpc_service_method.h
+++ b/include/grpc++/impl/rpc_service_method.h
@@ -167,7 +167,7 @@ class RpcServiceMethod : public RpcMethod {
                    MethodHandler* handler,
                    grpc::protobuf::Message* request_prototype,
                    grpc::protobuf::Message* response_prototype)
-      : RpcMethod(name, type),
+      : RpcMethod(name, type, nullptr),
         handler_(handler),
         request_prototype_(request_prototype),
         response_prototype_(response_prototype) {}
diff --git a/src/compiler/cpp_generator.cc b/src/compiler/cpp_generator.cc
index defb44aae9..ff3758ccfa 100644
--- a/src/compiler/cpp_generator.cc
+++ b/src/compiler/cpp_generator.cc
@@ -359,7 +359,7 @@ void PrintHeaderService(grpc::protobuf::io::Printer *printer,
       "class Stub GRPC_FINAL : public ::grpc::InternalStub {\n"
       " public:\n");
   printer->Indent();
-  printer->Print("Stub();\n");
+  printer->Print("Stub(const std::shared_ptr< ::grpc::ChannelInterface>& channel);\n");
   for (int i = 0; i < service->method_count(); ++i) {
     PrintHeaderClientMethod(printer, service->method(i), vars);
   }
@@ -744,15 +744,14 @@ void PrintSourceService(grpc::protobuf::io::Printer *printer,
       *vars,
       "std::unique_ptr< $ns$$Service$::Stub> $ns$$Service$::NewStub("
       "const std::shared_ptr< ::grpc::ChannelInterface>& channel) {\n"
-      "  std::unique_ptr< $ns$$Service$::Stub> stub(new $ns$$Service$::Stub());\n"
-      "  stub->set_channel(channel);\n"
+      "  std::unique_ptr< $ns$$Service$::Stub> stub(new $ns$$Service$::Stub(channel));\n"
       "  return stub;\n"
       "}\n\n");
-  printer->Print(*vars, "$ns$$Service$::Stub::Stub()");
+  printer->Print(*vars, "$ns$$Service$::Stub::Stub(const std::shared_ptr< ::grpc::ChannelInterface>& channel)\n");
   printer->Indent();
+  printer->Print(": ::grpc::InternalStub(channel)");
   for (int i = 0; i < service->method_count(); ++i) {
     const grpc::protobuf::MethodDescriptor *method = service->method(i);
-    (*vars)["Sep"] = (i==0) ? ":" : ",";
     (*vars)["Method"] = method->name();
     (*vars)["Idx"] = as_string(i);
     if (NoStreaming(method)) {
@@ -764,12 +763,13 @@ void PrintSourceService(grpc::protobuf::io::Printer *printer,
     } else {
       (*vars)["StreamingType"] = "BIDI_STREAMING";
     }
-    printer->Print(*vars, "$Sep$ rpcmethod_$Method$_("
+    printer->Print(*vars, ", rpcmethod_$Method$_("
       "$prefix$$Service$_method_names[$Idx$], "
-      "::grpc::RpcMethod::$StreamingType$"
+      "::grpc::RpcMethod::$StreamingType$, "
+      "channel->RegisterMethod($prefix$$Service$_method_names[$Idx$])"
       ")\n");
   }
-  printer->Print("{}\n");
+  printer->Print("{}\n\n");
   printer->Outdent();
 
   for (int i = 0; i < service->method_count(); ++i) {
diff --git a/src/cpp/client/generic_stub.cc b/src/cpp/client/generic_stub.cc
index 751f98f1b0..53084b33cc 100644
--- a/src/cpp/client/generic_stub.cc
+++ b/src/cpp/client/generic_stub.cc
@@ -43,7 +43,7 @@ std::unique_ptr<GenericClientAsyncReaderWriter> GenericStub::Call(
     CompletionQueue* cq, void* tag) {
   return std::unique_ptr<GenericClientAsyncReaderWriter>(
       new GenericClientAsyncReaderWriter(
-          channel_.get(), cq, RpcMethod(method.c_str(), RpcMethod::BIDI_STREAMING), context, tag));
+          channel_.get(), cq, RpcMethod(method.c_str(), RpcMethod::BIDI_STREAMING, nullptr), context, tag));
 }