From 0ab6db4fe88c14e8115190cc6682faca4391bfb1 Mon Sep 17 00:00:00 2001
From: Vijay Pai <vpai@google.com>
Date: Tue, 7 Jul 2020 13:49:11 -0700
Subject: [PATCH] Make sure that some ops don't start concurrently with Finish

---
 test/cpp/end2end/test_service_impl.cc | 27 +++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/test/cpp/end2end/test_service_impl.cc b/test/cpp/end2end/test_service_impl.cc
index a45f4cfe353..bd12d76fd51 100644
--- a/test/cpp/end2end/test_service_impl.cc
+++ b/test/cpp/end2end/test_service_impl.cc
@@ -502,12 +502,20 @@ CallbackTestServiceImpl::ResponseStream(
       if (num_msgs_sent_ == server_responses_to_send_ - 1 &&
           server_coalescing_api_ != 0) {
         num_msgs_sent_++;
-        StartWriteLast(&response_, WriteOptions());
+        {
+          std::lock_guard<std::mutex> l(finish_mu_);
+          if (!finished_) {
+            StartWriteLast(&response_, WriteOptions());
+          }
+        }
         // If we use WriteLast, we shouldn't wait before attempting Finish
         FinishOnce(Status::OK);
       } else {
         num_msgs_sent_++;
-        StartWrite(&response_);
+        std::lock_guard<std::mutex> l(finish_mu_);
+        if (!finished_) {
+          StartWrite(&response_);
+        }
       }
     }
     experimental::CallbackServerContext* const ctx_;
@@ -571,12 +579,15 @@ CallbackTestServiceImpl::BidiStream(
       if (ok) {
         num_msgs_read_++;
         response_.set_message(request_.message());
-        if (num_msgs_read_ == server_write_last_) {
-          StartWriteLast(&response_, WriteOptions());
-          // If we use WriteLast, we shouldn't wait before attempting Finish
-        } else {
-          StartWrite(&response_);
-          return;
+        std::lock_guard<std::mutex> l(finish_mu_);
+        if (!finished_) {
+          if (num_msgs_read_ == server_write_last_) {
+            StartWriteLast(&response_, WriteOptions());
+            // If we use WriteLast, we shouldn't wait before attempting Finish
+          } else {
+            StartWrite(&response_);
+            return;
+          }
         }
       }