Merge pull request #5052 from ctiller/compress

Fix race condition in transport API

1 files changed, 18 insertions, 2 deletions

diff --git a/include/grpc++/impl/codegen/sync_stream.h b/include/grpc++/impl/codegen/sync_stream.h
index 33d25e837c..9ae48bd23d 100644
--- a/include/grpc++/impl/codegen/sync_stream.h
+++ b/include/grpc++/impl/codegen/sync_stream.h
@@ -193,6 +193,15 @@ class ClientWriter : public ClientWriterInterface<W> {
     cq_.Pluck(&ops);
   }
 
+  void WaitForInitialMetadata() {
+    GPR_ASSERT(!context_->initial_metadata_received_);
+
+    CallOpSet<CallOpRecvInitialMetadata> ops;
+    ops.RecvInitialMetadata(context_);
+    call_.PerformOps(&ops);
+    cq_.Pluck(&ops);  // status ignored
+  }
+
   using WriterInterface<W>::Write;
   bool Write(const W& msg, const WriteOptions& options) GRPC_OVERRIDE {
     CallOpSet<CallOpSendMessage> ops;
@@ -213,6 +222,9 @@ class ClientWriter : public ClientWriterInterface<W> {
   /// Read the final response and wait for the final status.
   Status Finish() GRPC_OVERRIDE {
     Status status;
+    if (!context_->initial_metadata_received_) {
+      finish_ops_.RecvInitialMetadata(context_);
+    }
     finish_ops_.ClientRecvStatus(context_, &status);
     call_.PerformOps(&finish_ops_);
     GPR_ASSERT(cq_.Pluck(&finish_ops_));
@@ -221,7 +233,8 @@ class ClientWriter : public ClientWriterInterface<W> {
 
  private:
   ClientContext* context_;
-  CallOpSet<CallOpGenericRecvMessage, CallOpClientRecvStatus> finish_ops_;
+  CallOpSet<CallOpRecvInitialMetadata, CallOpGenericRecvMessage,
+            CallOpClientRecvStatus> finish_ops_;
   CompletionQueue cq_;
   Call call_;
 };
@@ -292,7 +305,10 @@ class ClientReaderWriter GRPC_FINAL : public ClientReaderWriterInterface<W, R> {
   }
 
   Status Finish() GRPC_OVERRIDE {
-    CallOpSet<CallOpClientRecvStatus> ops;
+    CallOpSet<CallOpRecvInitialMetadata, CallOpClientRecvStatus> ops;
+    if (!context_->initial_metadata_received_) {
+      ops.RecvInitialMetadata(context_);
+    }
     Status status;
     ops.ClientRecvStatus(context_, &status);
     call_.PerformOps(&ops);