Merge pull request #11758 from ctiller/write_completion

Write completion changes

2 files changed, 20 insertions, 2 deletions

diff --git a/test/cpp/end2end/async_end2end_test.cc b/test/cpp/end2end/async_end2end_test.cc
index e841a702d4..41090d161a 100644
--- a/test/cpp/end2end/async_end2end_test.cc
+++ b/test/cpp/end2end/async_end2end_test.cc
@@ -266,6 +266,7 @@ class AsyncEnd2endTest : public ::testing::TestWithParam<TestScenario> {
   }
 
   void TearDown() override {
+    gpr_tls_set(&g_is_async_end2end_test, 0);
     server_->Shutdown();
     void* ignored_tag;
     bool ignored_ok;
@@ -274,7 +275,6 @@ class AsyncEnd2endTest : public ::testing::TestWithParam<TestScenario> {
       ;
     stub_.reset();
     poll_overrider_.reset();
-    gpr_tls_set(&g_is_async_end2end_test, 0);
     grpc_recycle_unused_port(port_);
   }
 
@@ -396,6 +396,7 @@ TEST_P(AsyncEnd2endTest, WaitAndShutdownTest) {
   ResetStub();
   SendRpc(1);
   EXPECT_EQ(0, notify);
+  gpr_tls_set(&g_is_async_end2end_test, 0);
   server_->Shutdown();
   wait_thread.join();
   EXPECT_EQ(1, notify);
@@ -404,8 +405,9 @@ TEST_P(AsyncEnd2endTest, WaitAndShutdownTest) {
 TEST_P(AsyncEnd2endTest, ShutdownThenWait) {
   ResetStub();
   SendRpc(1);
-  server_->Shutdown();
+  std::thread t([this]() { server_->Shutdown(); });
   server_->Wait();
+  t.join();
 }
 
 // Test a simple RPC using the async version of Next
diff --git a/test/cpp/end2end/end2end_test.cc b/test/cpp/end2end/end2end_test.cc
index 1f4861a7e6..e54cd03ca2 100644
--- a/test/cpp/end2end/end2end_test.cc
+++ b/test/cpp/end2end/end2end_test.cc
@@ -757,6 +757,22 @@ TEST_P(End2endTest, RequestStreamTwoRequests) {
   EXPECT_TRUE(s.ok());
 }
 
+TEST_P(End2endTest, RequestStreamTwoRequestsWithWriteThrough) {
+  ResetStub();
+  EchoRequest request;
+  EchoResponse response;
+  ClientContext context;
+
+  auto stream = stub_->RequestStream(&context, &response);
+  request.set_message("hello");
+  EXPECT_TRUE(stream->Write(request, WriteOptions().set_write_through()));
+  EXPECT_TRUE(stream->Write(request, WriteOptions().set_write_through()));
+  stream->WritesDone();
+  Status s = stream->Finish();
+  EXPECT_EQ(response.message(), "hellohello");
+  EXPECT_TRUE(s.ok());
+}
+
 TEST_P(End2endTest, RequestStreamTwoRequestsWithCoalescingApi) {
   ResetStub();
   EchoRequest request;