diff --git a/include/grpc++/client_context.h b/include/grpc++/client_context.h index a6e8ccc67c4..91293d11231 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 ClientReader; +template class ClientWriter; +template class ClientReaderWriter; +template class ServerReader; +template class ServerWriter; +template class ServerReaderWriter; class ClientContext { public: @@ -71,7 +77,12 @@ class ClientContext { friend class CallOpBuffer; friend class Channel; - friend class StreamContext; + template friend class ::grpc::ClientReader; + template friend class ::grpc::ClientWriter; + template friend class ::grpc::ClientReaderWriter; + template friend class ::grpc::ServerReader; + template friend class ::grpc::ServerWriter; + template 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 metadata_; + std::multimap send_initial_metadata_; + std::multimap recv_initial_metadata_; + std::multimap trailing_metadata_; }; } // namespace grpc diff --git a/include/grpc++/stream.h b/include/grpc++/stream.h index 85a7261fb92..2791468490d 100644 --- a/include/grpc++/stream.h +++ b/include/grpc++/stream.h @@ -35,6 +35,7 @@ #define __GRPCPP_STREAM_H__ #include +#include #include #include #include @@ -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_; }; diff --git a/src/cpp/client/client_context.cc b/src/cpp/client/client_context.cc index 5c2772f5dfb..afacc81c168 100644 --- a/src/cpp/client/client_context.cc +++ b/src/cpp/client/client_context.cc @@ -72,7 +72,7 @@ system_clock::time_point ClientContext::absolute_deadline() { void ClientContext::AddMetadata(const grpc::string &meta_key, const grpc::string &meta_value) { - metadata_.insert(std::make_pair(meta_key, meta_value)); + send_initial_metadata_.insert(std::make_pair(meta_key, meta_value)); } void ClientContext::StartCancel() {} diff --git a/src/cpp/common/call.cc b/src/cpp/common/call.cc index 84791aa0a46..ffbc599d573 100644 --- a/src/cpp/common/call.cc +++ b/src/cpp/common/call.cc @@ -119,7 +119,7 @@ void CallOpBuffer::AddSendInitialMetadata( } void CallOpBuffer::AddSendInitialMetadata(ClientContext *ctx) { - AddSendInitialMetadata(&ctx->metadata_); + AddSendInitialMetadata(&ctx->send_initial_metadata_); } void CallOpBuffer::AddSendMessage(const google::protobuf::Message& message) {