Add client side WaitForInitialMetadata for streaming.

3 files changed, 44 insertions, 0 deletions

diff --git a/include/grpc++/impl/call.h b/include/grpc++/impl/call.h
index a1ef9268f0..5fafd0e890 100644
--- a/include/grpc++/impl/call.h
+++ b/include/grpc++/impl/call.h
@@ -134,7 +134,16 @@ class Call final {
   grpc_call *call() { return call_; }
   CompletionQueue *cq() { return cq_; }
 
+  // TODO(yangg) change it to a general state query function.
+  bool initial_metadata_received() {
+    return initial_metadata_received_;
+  }
+  void set_initial_metadata_received() {
+    initial_metadata_received_ = true;
+  }
+
  private:
+  bool initial_metadata_received_ = false;
   CallHook *call_hook_;
   CompletionQueue *cq_;
   grpc_call* call_;
diff --git a/include/grpc++/stream.h b/include/grpc++/stream.h
index a59ccd8c05..6da1be4e9f 100644
--- a/include/grpc++/stream.h
+++ b/include/grpc++/stream.h
@@ -98,7 +98,22 @@ class ClientReader final : public ClientStreamingInterface,
     cq_.Pluck(&buf);
   }
 
+  // Blocking wait for initial metadata from server. The received metadata
+  // can only be accessed after this call returns. Calling this method is
+  // optional as it will be called internally before the first Read.
+  void WaitForInitialMetadata() {
+    if (!call_.initial_metadata_received()) {
+      CallOpBuffer buf;
+      buf.AddRecvInitialMetadata(&context_->recv_initial_metadata_);
+      call_.PerformOps(&buf);
+      GPR_ASSERT(cq_.Pluck(&buf));
+      call_.set_initial_metadata_received();
+    }
+  }
+
+
   virtual bool Read(R *msg) override {
+    WaitForInitialMetadata();
     CallOpBuffer buf;
     bool got_message;
     buf.AddRecvMessage(msg, &got_message);
@@ -186,7 +201,21 @@ class ClientReaderWriter final : public ClientStreamingInterface,
     GPR_ASSERT(cq_.Pluck(&buf));
   }
 
+  // Blocking wait for initial metadata from server. The received metadata
+  // can only be accessed after this call returns. Calling this method is
+  // optional as it will be called internally before the first Read.
+  void WaitForInitialMetadata() {
+    if (!call_.initial_metadata_received()) {
+      CallOpBuffer buf;
+      buf.AddRecvInitialMetadata(&context_->recv_initial_metadata_);
+      call_.PerformOps(&buf);
+      GPR_ASSERT(cq_.Pluck(&buf));
+      call_.set_initial_metadata_received();
+    }
+  }
+
   virtual bool Read(R *msg) override {
+    WaitForInitialMetadata();
     CallOpBuffer buf;
     bool got_message;
     buf.AddRecvMessage(msg, &got_message);
diff --git a/src/cpp/common/call.cc b/src/cpp/common/call.cc
index a20d4a0d9a..aae69084eb 100644
--- a/src/cpp/common/call.cc
+++ b/src/cpp/common/call.cc
@@ -121,6 +121,12 @@ void CallOpBuffer::AddSendInitialMetadata(
   initial_metadata_ = FillMetadataArray(metadata);
 }
 
+void CallOpBuffer::AddRecvInitialMetadata(
+    std::multimap<grpc::string, grpc::string>* metadata) {
+  recv_initial_metadata_ = metadata;
+}
+
+
 void CallOpBuffer::AddSendInitialMetadata(ClientContext *ctx) {
   AddSendInitialMetadata(&ctx->send_initial_metadata_);
 }