3 files changed, 19 insertions, 10 deletions
diff --git a/src/cpp/client/client_unary_call.cc b/src/cpp/client/client_unary_call.cc
index 284af33b43..03a0326128 100644
--- a/src/cpp/client/client_unary_call.cc
+++ b/src/cpp/client/client_unary_call.cc
@@ -52,10 +52,10 @@ Status BlockingUnaryCall(ChannelInterface *channel, const RpcMethod &method,
   Status status;
   buf.AddSendInitialMetadata(context);
   buf.AddSendMessage(request);
-  buf.AddRecvInitialMetadata(&context->recv_initial_metadata_);
+  buf.AddRecvInitialMetadata(context);
   buf.AddRecvMessage(result);
   buf.AddClientSendClose();
-  buf.AddClientRecvStatus(&context->trailing_metadata_, &status);
+  buf.AddClientRecvStatus(context, &status);
   call.PerformOps(&buf);
   GPR_ASSERT((cq.Pluck(&buf) && buf.got_message) || !status.IsOk());
   return status;
@@ -79,10 +79,10 @@ void AsyncUnaryCall(ChannelInterface *channel, const RpcMethod &method,
   Call call(channel->CreateCall(method, context, cq));
   buf->AddSendInitialMetadata(context);
   buf->AddSendMessage(request);
-  buf->AddRecvInitialMetadata(&context->recv_initial_metadata_);
+  buf->AddRecvInitialMetadata(context);
   buf->AddRecvMessage(result);
   buf->AddClientSendClose();
-  buf->AddClientRecvStatus(&context->trailing_metadata_, status);
+  buf->AddClientRecvStatus(context, status);
   call.PerformOps(buf);
 }
 
diff --git a/src/cpp/common/call.cc b/src/cpp/common/call.cc
index f1142cf8e5..04af36f312 100644
--- a/src/cpp/common/call.cc
+++ b/src/cpp/common/call.cc
@@ -130,9 +130,9 @@ void CallOpBuffer::AddSendInitialMetadata(
   initial_metadata_ = FillMetadataArray(metadata);
 }
 
-void CallOpBuffer::AddRecvInitialMetadata(
-    std::multimap<grpc::string, grpc::string>* metadata) {
-  recv_initial_metadata_ = metadata;
+void CallOpBuffer::AddRecvInitialMetadata(ClientContext* ctx) {
+  ctx->initial_metadata_received_ = true;
+  recv_initial_metadata_ = &ctx->recv_initial_metadata_;
 }
 
 void CallOpBuffer::AddSendInitialMetadata(ClientContext* ctx) {
@@ -154,9 +154,8 @@ void CallOpBuffer::AddServerRecvClose(bool* cancelled) {
   recv_closed_ = cancelled;
 }
 
-void CallOpBuffer::AddClientRecvStatus(
-    std::multimap<grpc::string, grpc::string>* metadata, Status* status) {
-  recv_trailing_metadata_ = metadata;
+void CallOpBuffer::AddClientRecvStatus(ClientContext* context, Status* status) {
+  recv_trailing_metadata_ = &context->trailing_metadata_;
   recv_status_ = status;
 }
 
diff --git a/src/cpp/server/server_context.cc b/src/cpp/server/server_context.cc
index 21a61af3a0..df4c4dc314 100644
--- a/src/cpp/server/server_context.cc
+++ b/src/cpp/server/server_context.cc
@@ -57,4 +57,14 @@ ServerContext::~ServerContext() {
   }
 }
 
+void ServerContext::AddInitialMetadata(const grpc::string& key,
+                                  const grpc::string& value) {
+  initial_metadata_.insert(std::make_pair(key, value));
+}
+
+void ServerContext::AddTrailingMetadata(const grpc::string& key,
+                                  const grpc::string& value) {
+  trailing_metadata_.insert(std::make_pair(key, value));
+}
+
 }  // namespace grpc