Add trailing metadata to client context and use it.

2 files changed, 25 insertions, 8 deletions

diff --git a/include/grpc++/client_context.h b/include/grpc++/client_context.h
index a6e8ccc67c..91293d1123 100644
--- a/include/grpc++/client_context.h
+++ b/include/grpc++/client_context.h
@@ -50,6 +50,12 @@ struct grpc_completion_queue;
 namespace grpc {
 
 class CallOpBuffer;
+template <class R> class ClientReader;
+template <class W> class ClientWriter;
+template <class R, class W> class ClientReaderWriter;
+template <class R> class ServerReader;
+template <class W> class ServerWriter;
+template <class R, class W> class ServerReaderWriter;
 
 class ClientContext {
  public:
@@ -71,7 +77,12 @@ class ClientContext {
 
   friend class CallOpBuffer;
   friend class Channel;
-  friend class StreamContext;
+  template <class R> friend class ::grpc::ClientReader;
+  template <class W> friend class ::grpc::ClientWriter;
+  template <class R, class W> friend class ::grpc::ClientReaderWriter;
+  template <class R> friend class ::grpc::ServerReader;
+  template <class W> friend class ::grpc::ServerWriter;
+  template <class R, class W> friend class ::grpc::ServerReaderWriter;
 
   grpc_call *call() { return call_; }
   void set_call(grpc_call *call) {
@@ -87,7 +98,9 @@ class ClientContext {
   grpc_call *call_;
   grpc_completion_queue *cq_;
   gpr_timespec absolute_deadline_;
-  std::multimap<grpc::string, grpc::string> metadata_;
+  std::multimap<grpc::string, grpc::string> send_initial_metadata_;
+  std::multimap<grpc::string, grpc::string> recv_initial_metadata_;
+  std::multimap<grpc::string, grpc::string> trailing_metadata_;
 };
 
 }  // namespace grpc
diff --git a/include/grpc++/stream.h b/include/grpc++/stream.h
index 85a7261fb9..2791468490 100644
--- a/include/grpc++/stream.h
+++ b/include/grpc++/stream.h
@@ -35,6 +35,7 @@
 #define __GRPCPP_STREAM_H__
 
 #include <grpc++/channel_interface.h>
+#include <grpc++/client_context.h>
 #include <grpc++/completion_queue.h>
 #include <grpc++/impl/call.h>
 #include <grpc++/status.h>
@@ -87,7 +88,7 @@ class ClientReader final : public ClientStreamingInterface,
   ClientReader(ChannelInterface *channel, const RpcMethod &method,
                ClientContext *context,
                const google::protobuf::Message &request)
-      : call_(channel->CreateCall(method, context, &cq_)) {
+      : context_(context), call_(channel->CreateCall(method, context, &cq_)) {
     CallOpBuffer buf;
     buf.AddSendMessage(request);
     buf.AddClientSendClose();
@@ -105,13 +106,14 @@ class ClientReader final : public ClientStreamingInterface,
   virtual Status Finish() override {
     CallOpBuffer buf;
     Status status;
-    buf.AddClientRecvStatus(nullptr, &status); // TODO metadata
+    buf.AddClientRecvStatus(&context_->trailing_metadata_, &status);
     call_.PerformOps(&buf);
     GPR_ASSERT(cq_.Pluck(&buf));
     return status;
   }
 
  private:
+  ClientContext* context_;
   CompletionQueue cq_;
   Call call_;
 };
@@ -124,7 +126,7 @@ class ClientWriter final : public ClientStreamingInterface,
   ClientWriter(ChannelInterface *channel, const RpcMethod &method,
                ClientContext *context,
                google::protobuf::Message *response)
-      : response_(response),
+      : context_(context), response_(response),
         call_(channel->CreateCall(method, context, &cq_)) {}
 
   virtual bool Write(const W& msg) override {
@@ -146,13 +148,14 @@ class ClientWriter final : public ClientStreamingInterface,
     CallOpBuffer buf;
     Status status;
     buf.AddRecvMessage(response_);
-    buf.AddClientRecvStatus(nullptr, &status);  // TODO metadata
+    buf.AddClientRecvStatus(&context_->trailing_metadata_, &status);
     call_.PerformOps(&buf);
     GPR_ASSERT(cq_.Pluck(&buf));
     return status;
   }
 
  private:
+  ClientContext* context_;
   google::protobuf::Message *const response_;
   CompletionQueue cq_;
   Call call_;
@@ -167,7 +170,7 @@ class ClientReaderWriter final : public ClientStreamingInterface,
   // Blocking create a stream.
   ClientReaderWriter(ChannelInterface *channel,
                      const RpcMethod &method, ClientContext *context)
-      : call_(channel->CreateCall(method, context, &cq_)) {}
+      : context_(context), call_(channel->CreateCall(method, context, &cq_)) {}
 
   virtual bool Read(R *msg) override {
     CallOpBuffer buf;
@@ -193,13 +196,14 @@ class ClientReaderWriter final : public ClientStreamingInterface,
   virtual Status Finish() override {
     CallOpBuffer buf;
     Status status;
-    buf.AddClientRecvStatus(nullptr, &status);  // TODO metadata
+    buf.AddClientRecvStatus(&context_->trailing_metadata_, &status);
     call_.PerformOps(&buf);
     GPR_ASSERT(cq_.Pluck(&buf));
     return status;
   }
 
  private:
+  ClientContext* context_;
   CompletionQueue cq_;
   Call call_;
 };